mynode/rootfs/standard/usr/bin/mynode_usb_extras.py

#!/usr/bin/python3
import time
import os
import subprocess
import logging
import json
import atexit
from http.server import HTTPServer, SimpleHTTPRequestHandler
import pyudev
from systemd import journal
from threading import Thread
from utilities import *
from drive_info import *

log = logging.getLogger('mynode')
log.addHandler(journal.JournaldLogHandler())
log.setLevel(logging.INFO)
set_logger(log)

################################
## USB Device Cache
################################
usb_devices = []

def reset_usb_devices():
    global usb_devices
    for d in usb_devices:
        d.stop()
    usb_devices = []
    write_usb_devices_json()

def add_usb_device(usb_device):
    global usb_devices
    usb_devices.append(usb_device)
    write_usb_devices_json()

def remove_usb_device(usb_device):
    global usb_devices
    new_devices = []
    for d in usb_devices:
        if d.id != usb_device.id:
            new_devices.append(d)
    usb_devices = new_devices
    write_usb_devices_json()

def write_usb_devices_json():
    global usb_devices
    json_str = json.dumps([ob.to_dict() for ob in usb_devices])
    with open('/tmp/usb_extras.json', 'w') as f:
        f.write(json_str)
        f.close()


################################
## Utility Functions
################################
def set_usb_extras_state(state):
    log_message("USB Extras State: {}".format(state))
    try:
        with open("/tmp/.usb_extras_state", "w") as f:
            f.write(state)
        os.system("sync")
        return True
    except:
        return False
    return False


################################
## HTTP Server Functions
################################
class NoCacheHTTPRequestHandler(
    SimpleHTTPRequestHandler
):
    def send_response_only(self, code, message=None):
        super().send_response_only(code, message)
        self.send_header('Cache-Control', 'no-store, must-revalidate')
        self.send_header('Expires', '0')

def web_handler_from(directory):
    def _init(self, *args, **kwargs):
        return NoCacheHTTPRequestHandler.__init__(self, *args, directory=self.directory, **kwargs)
    return type(f'HandlerFrom<{directory}>',
                (NoCacheHTTPRequestHandler,),
                {'__init__': _init, 'directory': directory})

################################
## Detection Functions
################################
def check_partition_for_opendime(partition):
    is_opendime = False
    if mount_partition(partition, "temp_check"):
        if os.path.isfile("/mnt/usb_extras/temp_check/support/opendime.png"):
            is_opendime = True
    unmount_partition("temp_check")
    return is_opendime

################################
## Device Handlers
################################
usb_device_id = 0
                
class UsbDeviceHandler:
    def __init__(self):
        global usb_device_id
        self.id = usb_device_id
        usb_device_id = usb_device_id + 1

    def to_dict(self):
        raise NotImplementedError

class OpendimeHandler(UsbDeviceHandler):

    def __init__(self, block_device, partition):
        super().__init__()
        self.device = block_device
        self.device_type = "opendime"
        self.partition = partition
        self.folder_name = f"opendime_{self.id}"
        self.state = "loading_1"
        self.http_server = None
        self.http_server_thread = None

    def to_dict(self):
        dict = {}
        dict["id"] = self.id
        dict["device_type"] = self.device_type
        dict["device"] = self.device
        dict["partition"] = self.partition
        dict["folder_name"] = self.folder_name
        dict["port"] = self.port
        dict["state"] = self.state
        return dict

    def start(self):
        try:
            if mount_partition(self.partition, self.folder_name, "rw"):
                # Check device state
                self.state = "loading_2"
                try:
                    readme_file = f"/mnt/usb_extras/{self.folder_name}/README.txt"
                    private_key_file = f"/mnt/usb_extras/{self.folder_name}/private-key.txt"
                    if os.path.isfile(readme_file):
                        with open(readme_file) as f:
                            content = f.read()
                            if "This Opendime is fresh and unused. It hasn't picked a private key yet." in content:
                                self.state = "new"
                                log_message("  Opendime in state 'new'")

                    if os.path.isfile(private_key_file):
                        with open(private_key_file) as f:
                            content = f.read()
                            if "SEALED" in content:
                                self.state = "sealed"
                                log_message("  Opendime in state 'sealed'")
                            else:
                                self.state = "unsealed"
                                log_message("  Opendime in state 'unsealed'")

                except Exception as e:
                    self.state = "error_reading_opendime"

                self.port = 8010 + (self.id % 10)
                self.http_server = HTTPServer(('', self.port), web_handler_from(f"/mnt/usb_extras/{self.folder_name}"))
                self.http_server_thread = Thread(target = self.http_server.serve_forever)
                self.http_server_thread.setDaemon(True)
                self.http_server_thread.start()
                return True
            else:
                log_message("Error mounting partition for opendime")
                return False
        except Exception as e:
            unmount_partition(self.folder_name)
            log_message("Opendime Start Exception: {}".format(str(e)))
            return False

    def stop(self):
        try:
            if self.http_server:
                self.http_server.shutdown()
            unmount_partition(self.folder_name)
        except Exception as e:
            log_message("Opendime Stop Exception: {}".format(str(e)))

################################
## check_usb_devices()
################################
def check_usb_devices():
    try:
        # if new event, reset state
        # if no new event and in state (mounted), jump to state machine
        
        # Set initial state
        set_usb_extras_state("detecting")
        os.system("umount /mnt/usb_extras")

        # Detect drives
        drives = find_unmounted_drives()
        log_message(f"Drives: {drives}")

        # Check exactly one extra drive found
        drive_count = len(drives)
        if drive_count == 0:
            log_message("No USB extras found.")
        else:
            set_usb_extras_state("processing")
            for drive in drives:
                # Check drive for partitions
                drive = drives[0]
                partitions = find_partitions_for_drive(drive)
                log_message(f"Drive {drive} paritions: {partitions}")

                num_partitions = len(partitions)
                if num_partitions == 0:
                    log_message("No partitions found. Nothing to do.")
                elif num_partitions == 1:
                    # Process partition
                    partition = partitions[0]
                    log_message("One partition found! Scanning...")
                    if check_partition_for_opendime(partition):
                        log_message("Found Opendime!")
                        opendime = OpendimeHandler(drive, partition)
                        if opendime.start():
                            add_usb_device(opendime)
                        else:
                            opendime.stop()
                    else:
                        log_message(f"Drive {drive} could not be detected.")
                else:
                    log_message(f"{num_partitions} partitions found. Not sure what to do.")

        # Successful scan post init or usb action detected, mark homepage refresh
        os.system("touch /tmp/homepage_needs_refresh")
            
    except Exception as e:
        log_message("Exception: {}".format(str(e)))
        set_usb_extras_state("error")
        reset_usb_devices()
        log_message("Caught exception. Delaying 30s.")
        time.sleep(30)


################################
## Main
################################
def main():
    # Setup
    os.system("mkdir -p /mnt/usb_extras")

    # Start fresh and check USB devices once on startup 
    unmount_partition("*")
    reset_usb_devices()
    check_usb_devices()

    # Monitor USB and re-check on add/remove
    context = pyudev.Context()
    monitor = pyudev.Monitor.from_netlink(context)
    monitor.filter_by(subsystem='usb')
    # this is module level logger, can be ignored
    log_message("Starting to monitor for usb")
    monitor.start()
    log_message("Waiting on USB Event...")
    set_usb_extras_state("waiting")
    for device in iter(monitor.poll, None):
        log_message("")
        log_message("Got USB event: %s", device.action)
        if device.action == 'add':
            check_usb_devices()
        else:
            # HANDLE DEVICE REMOVAL BETTER? This resets all and re-scans
            reset_usb_devices()
            check_usb_devices()
        log_message("Waiting on USB Event...")
        set_usb_extras_state("waiting")
    

@atexit.register
def goodbye():
    log_message("ATEXIT: Resetting devices")
    unmount_partition("*")
    reset_usb_devices()
    log_message("ATEXIT: Done")

# This is the main entry point for the program
if __name__ == "__main__":
    while True:
        try:
            main()
        except Exception as e:
            set_usb_extras_state("error")
            log_message("Main Exception: {}".format(str(e)))
            log_message("Caught exception. Delaying 30s.")
            unmount_partition("*")
            reset_usb_devices()
            time.sleep(30)