From 3ed4e01c64bf2b7b01a22fde142621ea1a15b785 Mon Sep 17 00:00:00 2001 From: Mike Frysinger Date: Sat, 30 May 2015 20:24:33 +0800 Subject: [PATCH 1/6] trim trailing whitespace --- usbkill.py | 78 +++++++++++++++++++++++++++--------------------------- 1 file changed, 39 insertions(+), 39 deletions(-) diff --git a/usbkill.py b/usbkill.py index 1501b1d..b28d534 100644 --- a/usbkill.py +++ b/usbkill.py @@ -1,4 +1,4 @@ -# _ _ _ _ _ +# _ _ _ _ _ # | | | | (_) | | # _ _ ___| |__ | | _ _| | | # | | | |/___) _ \| |_/ ) | | | @@ -55,11 +55,11 @@ 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 +68,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 +79,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 +116,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 +136,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 +171,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 +192,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 +228,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 +244,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 +261,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,7 +276,7 @@ def loop(settings): for device in start_devices: if device not in current_devices: kill_computer(settings) - + sleep(settings['sleep_time']) def startup_checks(): @@ -290,22 +290,22 @@ def startup_checks(): # Check arguments args = sys.argv[1:] - - # Check for help + + # 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: 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") @@ -335,45 +335,45 @@ def startup_checks(): 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") - + os.remove(sources_path + "settings.ini") + # Load settings settings = load_settings(SETTINGS_FILE) settings['shut_down'] = 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__": # Run startup checks and load settings settings = startup_checks() - + # 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) From 692e6f3a032dbc1a004d8ee63b0f6adfd3175099 Mon Sep 17 00:00:00 2001 From: Mike Frysinger Date: Sat, 30 May 2015 20:29:43 +0800 Subject: [PATCH 2/6] unify settings paths into the one constant --- usbkill.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/usbkill.py b/usbkill.py index b28d534..2defa26 100644 --- a/usbkill.py +++ b/usbkill.py @@ -43,15 +43,18 @@ # Set the settings filename here SETTINGS_FILE = '/etc/usbkill/settings.ini' +SETTINGS_DIR = os.path.dirname(SETTINGS_FILE) help_message = """ 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'] @@ -323,15 +326,17 @@ 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") + 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 " + sources_path + "/\n") print("[NOTICE] Copying setting.ini to " + SETTINGS_FILE ) os.system("cp " + sources_path + "settings.ini " + SETTINGS_FILE) if not copy_settings: From 5067fe38e14084e385c3a1da1567ad8c2a0e9671 Mon Sep 17 00:00:00 2001 From: Mike Frysinger Date: Sat, 30 May 2015 20:29:58 +0800 Subject: [PATCH 3/6] make usbkill directly executable --- usbkill.py | 1 + 1 file changed, 1 insertion(+) mode change 100644 => 100755 usbkill.py diff --git a/usbkill.py b/usbkill.py old mode 100644 new mode 100755 index 2defa26..3434b5c --- a/usbkill.py +++ b/usbkill.py @@ -1,3 +1,4 @@ +#!/usr/bin/python # _ _ _ _ _ # | | | | (_) | | # _ _ ___| |__ | | _ _| | | From b44e0e6e2b483bf9152f8af7723d0df0066bf2cd Mon Sep 17 00:00:00 2001 From: Mike Frysinger Date: Sat, 30 May 2015 20:30:17 +0800 Subject: [PATCH 4/6] add a proper main() function --- usbkill.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/usbkill.py b/usbkill.py index 3434b5c..dd2c928 100755 --- a/usbkill.py +++ b/usbkill.py @@ -283,7 +283,7 @@ def loop(settings): sleep(settings['sleep_time']) -def startup_checks(): +def startup_checks(args): # Splash print(" _ _ _ _ _ \n" + " | | | | (_) | | \n" + @@ -292,9 +292,6 @@ 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) @@ -367,9 +364,9 @@ def startup_checks(): 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): @@ -383,3 +380,6 @@ def exit_handler(signum, frame): # Start main loop loop(settings) + +if __name__=="__main__": + main(sys.argv[1:]) From b8ad368dfeae297529db46e9882804a3661b5bcc Mon Sep 17 00:00:00 2001 From: Mike Frysinger Date: Sat, 30 May 2015 20:34:58 +0800 Subject: [PATCH 5/6] switch to argparse for command line processing --- usbkill.py | 33 +++++++++++++-------------------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/usbkill.py b/usbkill.py index dd2c928..59bf90c 100755 --- a/usbkill.py +++ b/usbkill.py @@ -25,6 +25,7 @@ __version__ = "1.0-rc.2" +import argparse import re import subprocess import platform @@ -46,7 +47,7 @@ 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. @@ -292,24 +293,16 @@ def startup_checks(args): " | |_| |___ | |_) ) _ (| | | | \n" + " |____/(___/|____/|_| \_)_|\_)_)\n") - # Check for help - if '-h' in args or '--help' in args: - sys.exit(help_message) + 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) - copy_settings = False - if '--cs' in args: - args.remove('--cs') - copy_settings = True - - shut_down = True - if '--no-shut-down' in 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. @@ -329,7 +322,7 @@ def startup_checks(args): # 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: + if not os.path.isfile(SETTINGS_FILE) or opts.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. " @@ -337,12 +330,12 @@ def startup_checks(args): + SETTINGS_DIR + " or in " + sources_path + "/\n") print("[NOTICE] Copying setting.ini to " + SETTINGS_FILE ) os.system("cp " + sources_path + "settings.ini " + SETTINGS_FILE) - if not copy_settings: + if not opts.copy_settings: os.remove(sources_path + "settings.ini") # 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']: From 1387b675cf981b87aa79d1fa98ea2b697817999e Mon Sep 17 00:00:00 2001 From: Mike Frysinger Date: Sat, 30 May 2015 20:39:11 +0800 Subject: [PATCH 6/6] use standard python commands for copying settings file --- usbkill.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/usbkill.py b/usbkill.py index 59bf90c..9bafef9 100755 --- a/usbkill.py +++ b/usbkill.py @@ -27,6 +27,7 @@ import argparse import re +import shutil import subprocess import platform import os, sys, signal @@ -323,15 +324,15 @@ def startup_checks(args): # 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 opts.copy_settings: - sources_path = os.path.dirname(os.path.realpath(__file__)) + '/' - if not os.path.isfile(sources_path + "settings.ini"): + 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 " + sources_path + "/\n") + + 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) + shutil.copyfile(sources_path, SETTINGS_FILE) if not opts.copy_settings: - os.remove(sources_path + "settings.ini") + os.remove(sources_path) # Load settings settings = load_settings(SETTINGS_FILE)