From 23c68c5458b53aa401ac11330128e40918297769 Mon Sep 17 00:00:00 2001
From: Bahtiar `kalkin-` Gadimov <bahtiar@gadimov.de>
Date: Fri, 17 Mar 2017 13:08:37 +0100
Subject: [PATCH] Add PersistentCollection helper to qubes.devices

---
 qubes/devices.py     | 50 +++++++++++++++++++++++++++++++++++++++++++-
 qubes/vm/__init__.py |  3 ++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/qubes/devices.py b/qubes/devices.py
index 9b0fe8c7..5b204b23 100644
--- a/qubes/devices.py
+++ b/qubes/devices.py
@@ -47,7 +47,6 @@ Such extension should provide:
 
 import qubes.utils
 
-
 class DeviceNotAttached(qubes.exc.QubesException, KeyError):
     '''Trying to detach not attached device'''
     pass
@@ -327,3 +326,52 @@ class BlockDevice(object):
         self.domain = domain
         self.devtype = devtype
 
+class PersistentCollection(object):
+
+    ''' Helper object managing persistent `DeviceAssignment`s.
+    '''
+
+    def __init__(self):
+        self._dict = {}
+
+    def add(self, assignment: DeviceAssignment):
+        ''' Add assignment to collection '''
+        assert assignment.persistent and assignment.frontend_domain
+        vm = assignment.backend_domain
+        ident = assignment.ident
+        key = (vm, ident)
+        assert key not in self._dict
+
+        self._dict[key] = assignment
+
+    def discard(self, assignment):
+        ''' Discard assignment from collection '''
+        assert assignment.persistent and assignment.frontend_domain
+        vm = assignment.backend_domain
+        ident = assignment.ident
+        key = (vm, ident)
+        if key not in self._dict:
+            raise KeyError
+        del self._dict[key]
+
+    def __contains__(self, device) -> bool:
+        vm = device.backend_domain
+        ident = device.ident
+        key = (vm, ident)
+        return key in self._dict
+
+    def get(self, device: DeviceInfo) -> DeviceAssignment:
+        ''' Returns the corresponding `qubes.devices.DeviceAssignment` for the
+            device. '''
+        vm = device.backend_domain
+        ident = device.ident
+        key = (vm, ident)
+        if key not in self._dict:
+            raise KeyError
+        return self._dict[key]
+
+    def __iter__(self):
+        return self._dict.values().__iter__()
+
+    def __len__(self) -> int:
+        return len(self._dict.keys())
diff --git a/qubes/vm/__init__.py b/qubes/vm/__init__.py
index 90b2362c..f2070ca0 100644
--- a/qubes/vm/__init__.py
+++ b/qubes/vm/__init__.py
@@ -250,10 +250,11 @@ class BaseVM(qubes.PropertyHolder, metaclass=BaseVMMeta):
         for devclass in self.devices:
             devices = lxml.etree.Element('devices')
             devices.set('class', devclass)
-            for device in self.devices[devclass].attached(persistent=True):
+            for device in self.devices[devclass].assignments(persistent=True):
                 node = lxml.etree.Element('device')
                 node.set('backend-domain', device.backend_domain.name)
                 node.set('id', device.ident)
+                node.set('options', device.options)
                 devices.append(node)
             element.append(devices)