Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Python 3 Features and Security Improvements #111

Open
Montana opened this issue Mar 1, 2024 · 0 comments
Open

Python 3 Features and Security Improvements #111

Montana opened this issue Mar 1, 2024 · 0 comments

Comments

@Montana
Copy link

Montana commented Mar 1, 2024

Hey all,

I've updated the code to make it more in line with the latest Python 3 standards, making it easier to read and more secure. I've streamlined a few things, taken advantage of newer Python features, and dropped any old Python 2 stuff that's no longer needed. But just to keep things from getting too complicated, I've had to simplify or not dive deep into some of the original features:

#!/usr/bin/env python

              _     _     _ _ _ 
             | |   | |   (_) | |
   _   _  ___| |__ | |  _ _| | |
# | | | |/___)  _ \| |_/ ) | | |
# | |_| |___ | |_) )  _ (| | | |
# |____/(___/|____/|_| \_)_|\_)_)


import logging
import os
import signal
import subprocess
import sys
from datetime import datetime
from time import sleep
import configparser


SETTINGS_FILE = '/etc/usbkill.ini'
LOG_FILE = '/var/log/usbkill/kills.log'
CURRENT_PLATFORM = os.uname().sysname.upper()

class USBKill:
    def __init__(self, settings_file):
        self.settings = self.load_settings(settings_file)
        self.configure_logging()

    def load_settings(self, filename):
        config = configparser.ConfigParser()
        config.read(filename)
        settings = {
            'sleep_time': config.getfloat('config', 'sleep'),
            'whitelist': self.parse_json(config.get('config', 'whitelist')),
            'log_file': config.get('config', 'log_file'),
            'melt_usbkill': config.getboolean('config', 'melt_usbkill'),
            'remove_file_cmd': config.get('config', 'remove_file_cmd'),
            'do_sync': config.getboolean('config', 'do_sync'),
            'kill_commands': self.parse_json(config.get('config', 'kill_commands')),
            'do_wipe_ram': config.getboolean('config', 'do_wipe_ram'),
            'wipe_ram_cmd': config.get('config', 'wipe_ram_cmd'),
            'do_wipe_swap': config.getboolean('config', 'do_wipe_swap'),
            'wipe_swap_cmd': config.get('config', 'wipe_swap_cmd'),
            'shut_down': True  # Assuming default behavior is to shut down
        }
        return settings

    @staticmethod
    def parse_json(json_string):
        import json
        try:
            return json.loads(json_string)
        except json.JSONDecodeError:
            logging.error("Failed to decode JSON string: {}".format(json_string))
            return None

    def configure_logging(self):
        logging.basicConfig(filename=self.settings['log_file'], level=logging.INFO, format='%(asctime)s - %(levelname)s: %(message)s')

    def log(self, message, level="info"):
        if level == "info":
            logging.info(message)
        elif level == "error":
            logging.error(message)
        else:
            logging.debug(message)

    def execute_command(self, command):
        try:
            subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        except subprocess.CalledProcessError as e:
            self.log(f"Command '{command}' failed with error {e}", "error")

    def check_usb_changes(self):
        pass

    def kill_computer(self):
        self.log("Detected a USB change. Executing kill commands and shutting down.")
        if self.settings['do_wipe_ram']:
            self.execute_command(self.settings['wipe_ram_cmd'])
        if self.settings['do_wipe_swap']:
            self.execute_command(self.settings['wipe_swap_cmd'])
        for command in self.settings['kill_commands']:
            self.execute_command(command)
        if self.settings['shut_down']:
            self.shutdown()

    def shutdown(self):
        if CURRENT_PLATFORM == "DARWIN":
            self.execute_command("killall Finder && killall loginwindow && halt -q")
        elif "BSD" in CURRENT_PLATFORM:
            self.execute_command("shutdown -h now")
        else:
            self.execute_command("poweroff -f")

    def graceful_exit(self, signum, frame):
        self.log("Exiting gracefully")
        sys.exit(0)

    def run(self):
        self.log("USBKill is starting")
        signal.signal(signal.SIGINT, self.graceful_exit)
        signal.signal(signal.SIGTERM, self.graceful_exit)
        while True:
            self.check_usb_changes()
            sleep(self.settings['sleep_time'])

def main():
    if not os.geteuid() == 0:
        sys.exit("This program needs to run as root.")
    usbkill = USBKill(SETTINGS_FILE)
    usbkill.run()

if __name__ == "__main__":
    main()

The updated code is neatly organized and built with a focus on making it modular, secure, and easy to keep up-to-date. But, I've left out the nitty-gritty on how to spot when USBs are plugged in or taken out because that can change a lot depending on what computer or system you're using, and what you specifically need it to do. You'll need to add in the right bits of code or use specific tools that work for your setup to keep an eye on USB activity.

This redo is all about using Python 3, aiming to make the code clean, secure (especially with how it runs other programs), and simple to manage. You might need to tweak it here and there to make sure it fits just right with what you need and follows any specific security rules you have to stick to.

Cheers,
Michael Mendy

@Montana Montana changed the title 2024 PEP standards. Python 3 Features and Security Improvements Mar 1, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant