From c41505e26c8be5e8fb76afbc31b9662540dfacf2 Mon Sep 17 00:00:00 2001
From: Alexander Rosenberg <zanderpkg@pm.me>
Date: Fri, 17 Nov 2023 21:00:33 -0800
Subject: [PATCH] Add ability to change default target in usbguard-menu.py

---
 usbguard-menu.py | 109 ++++++++++++++++++++++++++++++++++-------------
 1 file changed, 79 insertions(+), 30 deletions(-)

diff --git a/usbguard-menu.py b/usbguard-menu.py
index 7d46eee..c4d20ac 100755
--- a/usbguard-menu.py
+++ b/usbguard-menu.py
@@ -4,9 +4,16 @@ import subprocess
 import shlex
 
 class DeviceMode(Enum):
-    DENY = 0
+    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
@@ -16,7 +23,7 @@ class Device:
 
     def from_output_line(line):
         tokens =  shlex.split(line, False, True)
-        mode = DeviceMode.DENY
+        mode = DeviceMode.BLOCK
         id = None
         name = None
         hash = None
@@ -41,7 +48,7 @@ class Device:
             case DeviceMode.ALLOW:
                 mode = "Allowed"
             case _:
-                mode = "Denied"
+                mode = "Blocked"
         if len(self.name) == 0:
             name = "<no name>"
         else:
@@ -58,36 +65,78 @@ def run_fuzzel(entries, prompt=None):
                           text=True,
                           shell=False)
     if proc.returncode != 0:
-        return None
-    return entries[int(proc.stdout)]
+        return (-1, None)
+    index = int(proc.stdout)
+    return (index, entries[index])
 
-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))
+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
 
-device = run_fuzzel(devices)
+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)
 
-if device is None:
-    exit()
+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)
 
-action = run_fuzzel(['Allow', 'Deny', 'Reject'],
-                    device.name + ' > ')
+def update_device_mode(device):
+    (_, action) = run_fuzzel(['Allow', 'Block', 'Reject'],
+                             device.name + ' > ')
 
-if action is None:
-    exit()
+    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)
+    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()