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

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

    def __str__(self):
        match self:
            case DeviceMode.ALLOW:
                return 'Allow'
            case DeviceMode.BLOCK:
                return 'Block'

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.BLOCK
        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 = "Blocked"
        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 (-1, None)
    index = int(proc.stdout)
    return (index, entries[index])

def get_implicit_policy_target():
    proc = subprocess.run(['usbguard',
                           'get-parameter',
                           'ImplicitPolicyTarget'],
                          capture_output=True,
                          text=True,
                          check=True)
    match proc.stdout:
        case 'allow\n':
            return DeviceMode.ALLOW
        case 'block\n':
            return DeviceMode.BLOCK

def set_implicit_policy_target(target):
    match target:
        case DeviceMode.ALLOW:
            mode = 'allow'
        case DeviceMode.BLOCK:
            mode = 'block'
    subprocess.run(['usbguard',
                    'set-parameter',
                    'ImplicitPolicyTarget',
                    mode],
                   check=True)

def update_default_action():
    (index, _) = run_fuzzel(['Allow', 'Block'],
                            'Default action > ')
    match index:
        case 0:
            set_implicit_policy_target(DeviceMode.ALLOW)
        case 1:
            set_implicit_policy_target(DeviceMode.BLOCK)

def update_device_mode(device):
    (_, action) = run_fuzzel(['Allow', 'Block', '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)

def main():
    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))

    devices.append(f'Set default action (Current: {get_implicit_policy_target()})')
    (index, device) = run_fuzzel(devices)

    if index == len(devices) - 1:
        update_default_action()
    elif index != -1:
        update_device_mode(device)

main()