random-scripts/system-menu/usbguard-menu.py

143 lines
4.2 KiB
Python
Executable File

#!/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()