diff --git a/usbkill.py b/usbkill.py old mode 100644 new mode 100755 index 1501b1d..9bafef9 --- a/usbkill.py +++ b/usbkill.py @@ -1,4 +1,5 @@ -# _ _ _ _ _ +#!/usr/bin/python +# _ _ _ _ _ # | | | | (_) | | # _ _ ___| |__ | | _ _| | | # | | | |/___) _ \| |_/ ) | | | @@ -24,7 +25,9 @@ __version__ = "1.0-rc.2" +import argparse import re +import shutil import subprocess import platform import os, sys, signal @@ -43,23 +46,26 @@ # Set the settings filename here SETTINGS_FILE = '/etc/usbkill/settings.ini' +SETTINGS_DIR = os.path.dirname(SETTINGS_FILE) -help_message = """ +__doc__ = """ usbkill is a simple program with one goal: quickly shutdown the computer when a USB is inserted or removed. Events are logged in /var/log/usbkill/kills.log You can configure a whitelist of USB ids that are acceptable to insert and the remove. The USB id can be found by running the command 'lsusb'. -Settings can be changed in /etc/usbkill/settings +Settings can be changed in %(settings)s In order to be able to shutdown the computer, this program needs to run as root. -""" +""" % { + 'settings': SETTINGS_FILE, +} def log(settings, msg): log_file = settings['log_file'] - + contents = '\n{0} {1}\nCurrent state:\n'.format(str(datetime.now()), msg) with open(log_file, 'a+') as log: log.write(contents) - + # Log current USB state if CURRENT_PLATFORM.startswith("DARWIN"): os.system("system_profiler SPUSBDataType >> " + log_file) @@ -68,7 +74,7 @@ def log(settings, msg): def shred(settings): shredder = settings['remove_file_command'] - + # List logs and settings to be removed if settings['melt_usbkill']: settings['folders_to_remove'].append(os.path.dirname(settings['log_file'])) @@ -79,32 +85,32 @@ def shred(settings): else: settings['files_to_remove'].append(os.path.realpath(__file__)) settings['files_to_remove'].append(usbkill_folder + "/settings.ini") - + # Remove files and folders for _file in settings['files_to_remove'] + settings['folders_to_remove']: os.system(shredder + _file) - + def kill_computer(settings): # Log what is happening: if not settings['melt_usbkill']: # No need to spend time on logging if logs will be removed log(settings, "Detected a USB change. Dumping the list of connected devices and killing the computer...") - + # Shred as specified in settings shred(settings) - + if settings['do_sync']: # Sync the filesystem to save recent changes os.system("sync") - + # Execute kill commands in order. for command in settings['kill_commands']: os.system(command) - + else: # If syncing is risky because it might take too long, then sleep for 5ms. # This will still allow for syncing in most cases. sleep(0.05) - + if settings['shut_down']: # Use argument --no-shut-down to prevent a shutdown. # Finally poweroff computer immediately if CURRENT_PLATFORM.startswith("DARWIN"): @@ -116,7 +122,7 @@ def kill_computer(settings): else: # Linux-based systems - Will shutdown os.system("poweroff -f") - + # Exit the process to prevent executing twice (or more) all commands sys.exit(0) @@ -136,29 +142,29 @@ def check_inside(result, devices): try: result["Built-in_Device"] except KeyError: - + # Check if vendor_id/product_id is available for this one try: assert "vendor_id" in result and "product_id" in result # Append to the list of devices devices.append(DEVICE_RE[1].findall(result["vendor_id"])[0] + ':' + DEVICE_RE[1].findall(result["product_id"])[0]) except AssertionError: {} - + # Check if there is items inside try: # Looks like, do the while again for result_deep in result["_items"]: # Check what's inside the _items array check_inside(result_deep, devices) - + except KeyError: {} - + # Run the loop devices = [] for result in df[0]["_items"]: check_inside(result, devices) return devices - + def lsusb(): # A Python version of the command 'lsusb' that returns a list of connected usbids if CURRENT_PLATFORM.startswith("DARWIN"): @@ -171,9 +177,9 @@ def lsusb(): def program_present(program): if sys.version_info[0] == 3: # Python3 - from shutil import which + from shutil import which return which(program) != None - + else: """ Test if an executable exist in Python2 @@ -192,7 +198,7 @@ def is_exe(fpath): if is_exe(exe_file): return True return False - + def load_settings(filename): # Libraries that are only needed in this function: from json import loads as jsonloads @@ -228,7 +234,7 @@ def get_setting(name, gtype=''): # Read all lines of settings file config.read(filename) - + # Build settings settings = dict({ 'sleep_time' : get_setting('sleep', 'FLOAT'), @@ -244,14 +250,14 @@ def get_setting(name, gtype=''): }) return settings - + def loop(settings): # Main loop that checks every 'sleep_time' seconds if computer should be killed. # Allows only whitelisted usb devices to connect! # Does not allow usb device that was present during program start to disconnect! start_devices = lsusb() acceptable_devices = set(start_devices + settings['whitelist']) - + # Write to logs that loop is starting: msg = "[INFO] Started patrolling the USB ports every " + str(settings['sleep_time']) + " seconds..." log(settings, msg) @@ -261,7 +267,7 @@ def loop(settings): while True: # List the current usb devices current_devices = lsusb() - + # Check that no usbids are connected twice. # Two devices with same usbid implies a usbid copy attack if settings['double_usbid_detection'] and not len(current_devices) == len(set(current_devices)): @@ -276,10 +282,10 @@ def loop(settings): for device in start_devices: if device not in current_devices: kill_computer(settings) - + sleep(settings['sleep_time']) -def startup_checks(): +def startup_checks(args): # Splash print(" _ _ _ _ _ \n" + " | | | | (_) | | \n" + @@ -288,27 +294,16 @@ def startup_checks(): " | |_| |___ | |_) ) _ (| | | | \n" + " |____/(___/|____/|_| \_)_|\_)_)\n") - # Check arguments - args = sys.argv[1:] - - # Check for help - if '-h' in args or '--help' in args: - sys.exit(help_message) - - copy_settings = False - if '--cs' in args: - args.remove('--cs') - copy_settings = True - - shut_down = True - if '--no-shut-down' in args: + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('--cs', dest='copy_settings', default=False, action='store_true', + help='Copy settings to the ini file') + parser.add_argument('--no-shut-down', dest='shut_down', default=True, action='store_false', + help='Run all commands, but do not shut down') + opts = parser.parse_args(args) + + if not opts.shut_down: print("[NOTICE] Ready to execute all the (potentially destructive) commands, but NOT shut down the computer.") - args.remove('--no-shut-down') - shut_down = False - - # Check all other args - if len(args) > 0: - sys.exit("\n[ERROR] Argument not understood. Can only understand -h\n") # Check if program is run as root, else exit. # Root is needed to power off the computer. @@ -323,57 +318,62 @@ def startup_checks(): except subprocess.CalledProcessError: print("[NOTICE] FileVault is disabled. Sensitive data SHOULD be encrypted.") - if not os.path.isdir("/etc/usbkill/"): - os.mkdir("/etc/usbkill/") + if not os.path.isdir(SETTINGS_DIR): + os.mkdir(SETTINGS_DIR) # On first time use copy settings.ini to /etc/usebkill/settings.ini # If dev-mode, always copy and don't remove old settings - if not os.path.isfile(SETTINGS_FILE) or copy_settings: - sources_path = os.path.dirname(os.path.realpath(__file__)) + '/' - if not os.path.isfile(sources_path + "settings.ini"): - sys.exit("\n[ERROR] You have lost your settings file. Get a new copy of the settings.ini and place it in /etc/usbkill/ or in " + sources_path + "/\n") + if not os.path.isfile(SETTINGS_FILE) or opts.copy_settings: + sources_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'settings.ini') + if not os.path.isfile(sources_path): + sys.exit("\n[ERROR] You have lost your settings file. " + "Get a new copy of the settings.ini and place it in " + + SETTINGS_DIR + " or in " + os.path.join(sources_path) + "/\n") print("[NOTICE] Copying setting.ini to " + SETTINGS_FILE ) - os.system("cp " + sources_path + "settings.ini " + SETTINGS_FILE) - if not copy_settings: - os.remove(sources_path + "settings.ini") - + shutil.copyfile(sources_path, SETTINGS_FILE) + if not opts.copy_settings: + os.remove(sources_path) + # Load settings settings = load_settings(SETTINGS_FILE) - settings['shut_down'] = shut_down - + settings['shut_down'] = opts.shut_down + # Make sure no spaces a present in paths to be wiped. for name in settings['folders_to_remove'] + settings['files_to_remove']: if ' ' in name: msg += "[ERROR][WARNING] '" + name + "'as specified in your settings.ini contains a space.\n" sys.exit(msg) - + # Make sure srm is present if it will be used. if settings['melt_usbkill'] or len(settings['folders_to_remove'] + settings['files_to_remove']) > 0: if not program_present('srm'): sys.exit("[ERROR] usbkill configured to destroy data, but srm not installed.\n") if not settings['remove_file_command'].startswith('srm'): sys.exit("[ERROR] remove_file_command should start with `srm'. srm should be used for automated data overwrite.\n") - + # Make sure there is a logging folder log_folder = os.path.dirname(settings['log_file']) if not os.path.isdir(log_folder): os.mkdir(log_folder) - + return settings -if __name__=="__main__": +def main(args): # Run startup checks and load settings - settings = startup_checks() - + settings = startup_checks(args) + # Define exit handler now that settings are loaded... def exit_handler(signum, frame): print("\n[INFO] Exiting because exit signal was received\n") log(settings, "[INFO] Exiting because exit signal was received") sys.exit(0) - + # Register handlers for clean exit of program for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGQUIT, ]: signal.signal(sig, exit_handler) - + # Start main loop loop(settings) + +if __name__=="__main__": + main(sys.argv[1:])