From 2b1ebba547aa0b56f68e5faddd20ecb5a4e819a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Tue, 28 Feb 2017 01:34:09 +0100
Subject: [PATCH] vm: add actual VM classes

QubesOS/qubes-issues#853
---
 qubesmgmt/app.py         | 21 +++++++++++++++++++--
 qubesmgmt/vm/__init__.py | 28 ++++++++++++++++++++++++++--
 setup.py                 | 10 ++++++++++
 3 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/qubesmgmt/app.py b/qubesmgmt/app.py
index c856dd2..8f9c3d9 100644
--- a/qubesmgmt/app.py
+++ b/qubesmgmt/app.py
@@ -29,6 +29,7 @@ import subprocess
 import qubesmgmt.base
 import qubesmgmt.vm
 import qubesmgmt.exc
+import qubesmgmt.utils
 
 QUBESD_SOCK = '/var/run/qubesd.sock'
 BUF_SIZE = 4096
@@ -39,6 +40,7 @@ class VMCollection(object):
     def __init__(self, app):
         self.app = app
         self._vm_list = None
+        self._vm_objects = {}
 
     def clear_cache(self):
         '''Clear cached list of VMs'''
@@ -61,12 +63,27 @@ class VMCollection(object):
                 [vm_prop.split('=', 1) for vm_prop in props])
 
         self._vm_list = new_vm_list
+        for name, vm in self._vm_objects.items():
+            if vm.name not in self._vm_list:
+                # VM no longer exists
+                del self._vm_objects[name]
+            elif vm.__class__.__name__ != self._vm_list[vm.name]['class']:
+                # VM class have changed
+                del self._vm_objects[name]
+            # TODO: some generation ID, to detect VM re-creation
+            elif name != vm.name:
+                # renamed
+                self._vm_objects[vm.name] = vm
+                del self._vm_objects[name]
 
     def __getitem__(self, item):
         if item not in self:
             raise KeyError(item)
-        return qubesmgmt.vm.QubesVM(
-            self.app, item, self._vm_list[item]['class'])
+        if item not in self._vm_objects:
+            cls = qubesmgmt.utils.get_entry_point_one('qubesmgmt.vm',
+                self._vm_list[item]['class'])
+            self._vm_objects[item] = cls(self.app, item)
+        return self._vm_objects[item]
 
     def __contains__(self, item):
         self.refresh_cache()
diff --git a/qubesmgmt/vm/__init__.py b/qubesmgmt/vm/__init__.py
index 498858a..51bac49 100644
--- a/qubesmgmt/vm/__init__.py
+++ b/qubesmgmt/vm/__init__.py
@@ -25,8 +25,7 @@ import qubesmgmt.base
 
 class QubesVM(qubesmgmt.base.PropertyHolder):
     '''Qubes domain.'''
-    def __init__(self, app, name, vm_class):
-        self._class = vm_class
+    def __init__(self, app, name):
         super(QubesVM, self).__init__(app, 'mgmt.vm.property.', name)
 
     @property
@@ -112,3 +111,28 @@ class QubesVM(qubesmgmt.base.PropertyHolder):
         '''
         raise NotImplementedError
         #self.qubesd_call(self._method_dest, 'mgmt.vm.Resume')
+
+
+class AdminVM(QubesVM):
+    '''Dom0'''
+    pass
+
+
+class AppVM(QubesVM):
+    '''Application VM'''
+    pass
+
+
+class StandaloneVM(QubesVM):
+    '''Standalone Application VM'''
+    pass
+
+
+class TemplateVM(QubesVM):
+    '''Template for AppVM'''
+    pass
+
+
+class DispVM(QubesVM):
+    '''Disposable VM'''
+    pass
diff --git a/setup.py b/setup.py
index 98b7c77..d86d851 100644
--- a/setup.py
+++ b/setup.py
@@ -12,4 +12,14 @@ if __name__ == '__main__':
         license='LGPL2.1+',
         url='https://www.qubes-os.org/',
         packages=setuptools.find_packages(),
+        entry_points={
+            'qubesmgmt.vm': [
+                'AppVM = qubesmgmt.vm:AppVM',
+                'TemplateVM = qubesmgmt.vm:TemplateVM',
+                'StandaloneVM = qubesmgmt.vm:StandaloneVM',
+                'AdminVM = qubesmgmt.vm:AdminVM',
+                'DispVM = qubesmgmt.vm:DispVM',
+            ],
+        },
+
         )