#!/usr/bin/env python
from enum import Enum
import subprocess
import shlex

class DeviceMode(Enum):
    DENY = 0
    ALLOW = 1

class Device:
    def __init__(self, mode, id, name, hash):
        self.mode = mode
        self.id = id
        self.name = name
        self.hash =  hash

    def from_output_line(line):
        tokens =  shlex.split(line, False, True)
        mode = DeviceMode.DENY
        id = None
        name = None
        hash = None
        for i, token in enumerate(tokens):
            if i == len(tokens) - 1:
                next = None
            else:
                next = tokens[i + 1]
            match token:
                case 'allow':
                    mode =  DeviceMode.ALLOW
                case 'id':
                    id = next
                case 'name':
                    name = next
                case 'hash':
                    hash = next
        return Device(mode, id, name, hash)

    def __str__(self):
        match self.mode:
            case DeviceMode.ALLOW:
                mode = "Allowed"
            case _:
                mode = "Denied"
        if len(self.name) == 0:
            name = "<no name>"
        else:
            name =  self.name
        return f'{self.id}: {name} ({mode}, Hash: {self.hash})'

def run_fuzzel(entries, prompt=None):
    cmd = ["fuzzel", "-d", "--index"]
    if prompt is not None:
        cmd.append(f'--prompt={prompt}')
    proc = subprocess.run(cmd,
                          input='\n'.join(str(e) for e in entries),
                          capture_output=True,
                          text=True,
                          shell=False)
    if proc.returncode != 0:
        return None
    return entries[int(proc.stdout)]

devices = []
proc = subprocess.run(["usbguard", "list-devices"],
                      capture_output=True, text=True, check=True)
for line in proc.stdout.splitlines():
    devices.append(Device.from_output_line(line))

device = run_fuzzel(devices)

if device is None:
    exit()

action = run_fuzzel(['Allow', 'Deny', 'Reject'],
                    device.name + ' > ')

if action is None:
    exit()

match action:
    case 'Allow':
        subprocess.check_output(['usbguard',
                                 'allow-device',
                                 device.id], shell=False)
    case 'Reject':
        subprocess.check_output(['usbguard',
                                 'reject-device',
                                 device.id], shell=False)
    case _:
        subprocess.check_output(['usbguard',
                                 'block-device',
                                 device.id], shell=False)