From 0f78d257d49277fa461bbaf93741071b20e481c0 Mon Sep 17 00:00:00 2001
From: Bahtiar `kalkin-` Gadimov <bahtiar@gadimov.de>
Date: Fri, 7 Apr 2017 20:26:05 +0200
Subject: [PATCH] qvm-device add support for assignments

- Add header field ASSIGNED
- Persistent assignments are displayed with an asterisk
- Add '-p' to attach command

Fix qvm-device

Signed-off-by: Bahtiar `kalkin-` Gadimov <bahtiar@gadimov.de>
---
 qubes/tools/qvm_device.py | 74 +++++++++++++++++++++++++++++++--------
 1 file changed, 59 insertions(+), 15 deletions(-)

diff --git a/qubes/tools/qvm_device.py b/qubes/tools/qvm_device.py
index 90bf0377..a51972fd 100644
--- a/qubes/tools/qvm_device.py
+++ b/qubes/tools/qvm_device.py
@@ -49,17 +49,30 @@ def prepare_table(dev_list):
     output = []
     header = []
     if sys.stdout.isatty():
-        header += [('BACKEND:DEVID', 'DESCRIPTION', 'USED BY')]  # NOQA
+        header += [('VMNAME:DEVID', 'DESCRIPTION', 'USED BY', 'ASSIGNED')]  # NOQA
 
     for dev in dev_list:
         output += [(
-            "{!s}:{!s}".format(dev.backend_domain, dev.ident),
+            dev.id,
             dev.description,
-            str(dev.frontend_domain) if dev.frontend_domain else "",
+            str(dev.attached_to),
+            dev.assignments
         )]
 
     return header + sorted(output)
 
+class Line(object):
+
+    def __init__(self, device: qubes.devices.DeviceInfo, attached_to = None):
+        self.id = "{!s}:{!s}".format(device.backend_domain, device.ident)
+        self.description = device.description
+        self.attached_to = attached_to if attached_to else ""
+        self.frontends = []
+
+    @property
+    def assignments(self):
+        return ', '.join(self.frontends)
+
 
 def list_devices(args):
     ''' Called by the parser to execute the qubes-devices list
@@ -67,32 +80,60 @@ def list_devices(args):
     app = args.app
 
     result = []
+    devices = set()
     if hasattr(args, 'domains') and args.domains:
         for domain in args.domains:
-            result.extend(domain.devices[args.devclass].attached())
-    else:
-        for backend in app.domains:
-            result.extend(backend.devices[args.devclass])
+            for dev in domain.devices[args.devclass].attached():
+                devices.add(dev)
+            for dev in domain.devices[args.devclass].available():
+                devices.add(dev)
 
-    qubes.tools.print_table(prepare_table(result))
+    else:
+        for domain in app.domains:
+            for dev in domain.devices[args.devclass].available():
+                devices.add(dev)
+
+    result = {dev: Line(dev) for dev in devices}
+
+    for dev in result:
+        for domain in app.domains:
+            if domain == dev.backend_domain:
+                continue
+            elif dev in domain.devices[args.devclass].attached():
+                result[dev].attached_to = str(domain)
+
+            if dev in domain.devices[args.devclass].assignments():
+                if dev in domain.devices[args.devclass].persistent():
+                    result[dev].frontends.append(str(domain))
+
+
+    qubes.tools.print_table(prepare_table(result.values()))
 
 
 def attach_device(args):
     ''' Called by the parser to execute the :program:`qvm-devices attach`
         subcommand.
     '''
-    device = args.device
+    device_assignment = args.device_assignment
     vm = args.domains[0]
-    vm.devices[args.devclass].attach(device)
+    app = args.app
+    device_assignment.persistent = args.persistent
+    vm.devices[args.devclass].attach(device_assignment)
+    if device_assignment.persistent:
+        app.save()
 
 
 def detach_device(args):
     ''' Called by the parser to execute the :program:`qvm-devices detach`
         subcommand.
     '''
-    device = args.device
+    device_assignment = args.device_assignment
     vm = args.domains[0]
-    vm.devices[args.devclass].detach(device)
+    before = len(vm.devices[args.devclass].persistent())
+    vm.devices[args.devclass].detach(device_assignment)
+    after = len(vm.devices[args.devclass].persistent())
+    if after < before:
+        args.app.save()
 
 
 def init_list_parser(sub_parsers):
@@ -141,7 +182,7 @@ class DeviceAction(qubes.tools.QubesAction):
                 parser.error_runtime("no backend vm {!r}".format(vmname))
 
             try:
-                device = vm.devices[devclass][device_id]
+                vm.devices[devclass][device_id]
             except KeyError:
                 parser.error_runtime(
                     "backend vm {!r} doesn't expose device {!r}"
@@ -176,8 +217,8 @@ def get_parser(device_class=None):
     detach_parser = sub_parsers.add_parser(
         "detach", help="Detach device from domain", aliases=('d', 'dt'))
 
-    attach_parser.add_argument('VMNAME', action=qubes.tools.RunningVmNameAction)
-    detach_parser.add_argument('VMNAME', action=qubes.tools.RunningVmNameAction)
+    attach_parser.add_argument('VMNAME', action=qubes.tools.VmNameAction)
+    detach_parser.add_argument('VMNAME', action=qubes.tools.VmNameAction)
 
     if device_class == 'block':
         attach_parser.add_argument(metavar='BACKEND:DEVICE_ID', dest='device',
@@ -188,6 +229,9 @@ def get_parser(device_class=None):
         attach_parser.add_argument(metavar='BACKEND:DEVICE_ID',
                                    dest='device',
                                    action=DeviceAction)
+        attach_parser.add_argument('-p', '--persistent', default=False,
+                                   help='device will attached on each start of the VMNAME',
+                                   action='store_true')
         detach_parser.add_argument(metavar='BACKEND:DEVICE_ID',
                                    dest='device',
                                    action=DeviceAction)