#!/bin/bash
set -e

VERSION="1.2"
SYSCTL_VAR=kernel.deny_new_usb
NEW_DEVICES=
OLD_DEVICES=
TEMPORARY_WAIT=60
SUDO="sudo"

if [[ ${EUID} -eq 0 ]]; then
	SUDO=""
fi

usage() {
	echo "Usage: $(basename "$0") [COMMAND]"
	echo "Control usb device and protection settings."
	echo
	echo "COMMANDS:"
	echo "  protect, disable, off -- disallow new usb devices (protected)"
	echo "  unprotect, enable, on -- allow new usb devices (unprotected)"
	echo "  temporary, temp, tmp  -- temporarily disable protection (default ${TEMPORARY_WAIT} sec)"
	echo "  check                 -- exit with 1 if usb is unprotected"
	echo "  status                -- print current protection status"
	echo "  list, ls              -- list currently connected usb devices"
	echo "  log                   -- display usb events in the kernel ring buffer"
	echo "  version               -- display version information and exit"
}

version() {
	echo "$(basename "$0") ${VERSION}"
}

deny_new_usb() {
	${SUDO} sysctl -q "${SYSCTL_VAR}=${1}"
	usb_status
}

usb_protected() {
	sysctl -n "${SYSCTL_VAR}"|grep -q 1
}

usb_status() {
	if usb_protected; then
		printf '\e[32;1m'
	else
		printf '\e[31;1mUN'
	fi
	printf 'PROTECTED\e[0m\n'
}

usb_list() {
	lsusb -t
}

temporary() {
	OLD_DEVICES=$(usb_list)
	deny_new_usb 0
	echo 'Press Ctrl-C to finish'
	trap temporary_end INT TERM
	for _ in $(seq ${TEMPORARY_WAIT}); do
		sleep 1
		NEW_DEVICES=$(usb_list)
		diff_usb_list "${OLD_DEVICES}" "${NEW_DEVICES}"
		OLD_DEVICES="${NEW_DEVICES}"
	done
	temporary_end
}

temporary_end() {
	deny_new_usb 1
	diff_usb_list "${OLD_DEVICES}" "$(usb_list)"
	exit
}

diff_usb_list() {
	local OLD_DEVICES="${1}"
	local NEW_DEVICES="${2}"
	diff --color <(printf "%s" "${OLD_DEVICES}") <(printf "%s" "${NEW_DEVICES}") ||:
}

usb_log() {
	DMESG=dmesg
	if sysctl -n kernel.dmesg_restrict | grep -q 1; then
		DMESG="${SUDO} ${DMESG}"
	fi
	${DMESG} --color="${USBCTL_LOG_COLOR:-always}" --ctime|grep -i usb --color="${USBCTL_GREP_COLOR:-never}"
}

case "$1" in
	on|enable|unprotect)
		deny_new_usb 0
		;;
	off|disable|protect)
		deny_new_usb 1
		;;
	temporary|tmp|temp)
		temporary
		;;
	check)
		usb_protected
		;;
	status)
		usb_status
		;;
	list|ls)
		usb_list
		;;
	log)
		usb_log
		;;
	version|--version)
		version
		;;
	*)
		usage
		;;
esac

# vim:set noet:
