Show "List USB Devices" option if sys-usb exists

Also show individual menu items if user has more than one usbvm

resolves: QubesOS/qubes-issues#10131
resolves: QubesOS/qubes-issues#10202

Author: alimirjamali

qui/devices/actionable_widgets.py

@@ -361,20 +361,20 @@ async def widget_action(self, *_args):
         )
 
 
-#### Start sys-usb
+#### Start sys-usb (or other custom USBVMs)
 
 
-class StartSysUsb(ActionableWidget, SimpleActionWidget):
-    def __init__(self, sysusb: backend.VM, variant: str = "dark"):
+class StartUSBVM(ActionableWidget, SimpleActionWidget):
+    def __init__(self, usbvm: backend.VM, variant: str = "dark"):
         super().__init__(
-            icon_name=sysusb.icon_name,
-            text="<b>List USB Devices " "(start sys-usb)</b>",
+            icon_name=usbvm.icon_name,
+            text="<b>List USB Devices " "(start {})</b>".format(usbvm.name),
             variant=variant,
         )
-        self.sysusb = sysusb
+        self.usbvm = usbvm
 
     async def widget_action(self, *_args):
-        self.sysusb.vm_object.start()
+        self.usbvm.vm_object.start()
 
 
 #### Configuration-related actions

qui/devices/device_widget.py

@@ -36,6 +36,8 @@
 import qubesadmin.tests
 import qubesadmin.tests.mock_app
 
+from qubesadmin.device_protocol import DeviceCategory
+
 import qui
 import qui.utils
 
@@ -68,6 +70,15 @@
 DEV_TYPES = ["block", "usb", "mic"]
 
 
+def _is_usbvm(vm: qubesadmin.vm.QubesVM) -> bool:
+    """Helper function to detect USBVMS"""
+    for dev in vm.devices["pci"].get_assigned_devices():
+        for interface in dev.device.interfaces:
+            if interface.category == DeviceCategory.PCI_USB:
+                return True
+    return False
+
+
 class DeviceMenu(Gtk.Menu):
     """Menu for handling a single device"""
 
@@ -112,7 +123,7 @@ def __init__(self, app_name, qapp, dispatcher):
         self.vms: Set[backend.VM] = set()
         self.dispvm_templates: Set[backend.VM] = set()
         self.parent_ports_to_hide = []
-        self.sysusb: backend.VM | None = None
+        self.dormant_usbvms: Set[backend.VM] = set()
         self.dev_update_queue: Set = set()
         self.vm_update_queue: Set = set()
 
@@ -144,6 +155,10 @@ def __init__(self, app_name, qapp, dispatcher):
                 "device-unassign:" + devclass, self.device_unassigned
             )
 
+        self.dispatcher.add_handler("device-assign:pci" , self.pci_assigned)
+        self.dispatcher.add_handler("device-unassign:pci", self.pci_unassigned)
+        self.dispatcher.add_handler("domain-delete", self.remove_domain_item)
+
         self.dispatcher.add_handler("domain-shutdown", self.vm_shutdown)
         self.dispatcher.add_handler("domain-start-failed", self.vm_shutdown)
         self.dispatcher.add_handler("domain-start", self.vm_start)
@@ -242,9 +257,8 @@ def initialize_vm_data(self):
                         self.vms.add(wrapped_vm)
                     if wrapped_vm.is_dispvm_template:
                         self.dispvm_templates.add(wrapped_vm)
-                if vm.name == "sys-usb":
-                    self.sysusb = wrapped_vm
-                    self.sysusb.is_running = vm.is_running()
+                if not vm.is_running() and _is_usbvm(vm):
+                    self.dormant_usbvms.add(wrapped_vm)
             except qubesadmin.exc.QubesException:
                 # we don't have access to VM state
                 pass
@@ -360,6 +374,23 @@ def device_unassigned(self, vm, _event, device, **_kwargs):
             # assigned. Cheers!
             return
 
+    def pci_assigned(self, vm, _event, device, **_kwargs):
+        if not vm.is_running() and _is_usbvm(vm):
+            wrapped_vm = backend.VM(vm)
+            self.dormant_usbvms.add(wrapped_vm)
+
+    def pci_unassigned(self, vm, _event, device, **_kwargs):
+        wrapped_vm = backend.VM(vm)
+        if wrapped_vm in self.dormant_usbvms and not _is_usbvm(vm):
+            self.dormant_usbvms.discard(wrapped_vm)
+
+    def remove_domain_item(self, _submitter, _event, vm, **_kwargs):
+        # In a pefect world, core should trigger `device-unassign:pci` event and
+        # this method should not be necessary. But we are not certain :/
+        wrapped_vm = backend.VM(vm)
+        if wrapped_vm in self.dormant_usbvms:
+            self.dormant_usbvms.discard(wrapped_vm)
+
     def update_single_feature(self, _vm, _event, feature, value=None, oldvalue=None):
         if not value:
             new = set()
@@ -533,8 +564,8 @@ def vm_start(self, vm, _event, **_kwargs):
             internal, attachable = False, False
         if attachable and not internal:
             self.vms.add(wrapped_vm)
-        if wrapped_vm == self.sysusb:
-            self.sysusb.is_running = True
+        if wrapped_vm in self.dormant_usbvms:
+            self.dormant_usbvms.discard(wrapped_vm)
 
         for devclass in DEV_TYPES:
             try:
@@ -548,8 +579,8 @@ def vm_start(self, vm, _event, **_kwargs):
 
     def vm_shutdown(self, vm, _event, **_kwargs):
         wrapped_vm = backend.VM(vm)
-        if wrapped_vm == self.sysusb:
-            self.sysusb.is_running = False
+        if _is_usbvm(vm):
+            self.dormant_usbvms.add(wrapped_vm)
 
         self.vms.discard(wrapped_vm)
         self.dispvm_templates.discard(wrapped_vm)
@@ -633,13 +664,13 @@ def show_menu(self, _unused, _event):
 
             menu_items.append(device_item)
 
-        if not self.sysusb.is_running:
-            sysusb_item = actionable_widgets.generate_wrapper_widget(
+        for wrapped_vm in self.dormant_usbvms:
+            usbvm_item = actionable_widgets.generate_wrapper_widget(
                 Gtk.MenuItem,
                 "activate",
-                actionable_widgets.StartSysUsb(self.sysusb, theme),
+                actionable_widgets.StartUSBVM(wrapped_vm, theme),
             )
-            menu_items.append(sysusb_item)
+            menu_items.append(usbvm_item)
 
         for item in menu_items:
             tray_menu.add(item)