From 9bb59cdd20fa551563a04346559be2f09500708a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Sun, 6 Aug 2017 12:22:47 +0200
Subject: [PATCH] vm: add DispVMWrapper for calling a single service in new
 DispVM

This is a wrapper to use `$dispvm` target of qrexec call, just like any
other service call in qubesadmin module - using vm.run_service().
When running in dom0, qrexec-client-vm is not available, so DispVM needs
to be created "manually", using appropriate Admin API call
(admin.vm.CreateDisposable).

QubesOS/qubes-issues#2974
---
 qubesadmin/vm/__init__.py | 45 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

diff --git a/qubesadmin/vm/__init__.py b/qubesadmin/vm/__init__.py
index e131dba..a239640 100644
--- a/qubesadmin/vm/__init__.py
+++ b/qubesadmin/vm/__init__.py
@@ -346,7 +346,50 @@ class TemplateVM(QubesVM):
             except AttributeError:
                 pass
 
+class DispVMWrapper(QubesVM):
+    '''Wrapper class for new DispVM, supporting only service call
+
+    Note that when running in dom0, one need to manually kill the DispVM after
+    service call ends.
+    '''
+
+    def run_service(self, service, **kwargs):
+        if self.app.qubesd_connection_type == 'socket':
+            # create dispvm at service call
+            if self._method_dest.startswith('$dispvm'):
+                if self._method_dest.startswith('$dispvm:'):
+                    method_dest = self._method_dest[len('$dispvm:'):]
+                else:
+                    method_dest = 'dom0'
+                dispvm = self.app.qubesd_call(method_dest,
+                    'admin.vm.CreateDisposable')
+                dispvm = dispvm.decode('ascii')
+                self._method_dest = dispvm
+        return super(DispVMWrapper, self).run_service(service, **kwargs)
+
+    def cleanup(self):
+        '''Cleanup after DispVM usage'''
+        # in 'remote' case nothing is needed, as DispVM is cleaned up
+        # automatically
+        if self.app.qubesd_connection_type == 'socket':
+            try:
+                self.kill()
+            except qubesadmin.exc.QubesVMNotRunningError:
+                pass
+
 
 class DispVM(QubesVM):
     '''Disposable VM'''
-    pass
+
+    @classmethod
+    def from_appvm(cls, app, appvm):
+        '''Returns a wrapper for calling service in a new DispVM based on given
+        AppVM. If *appvm* is none, use default DispVM template'''
+
+        if appvm:
+            method_dest = '$dispvm:' + str(appvm)
+        else:
+            method_dest = '$dispvm'
+
+        wrapper = DispVMWrapper(app, method_dest)
+        return wrapper