From c83deccdd39005e64666f08040f83f3f2a178720 Mon Sep 17 00:00:00 2001
From: Rusty Bird <rustybird@net-c.com>
Date: Sun, 14 Jan 2018 18:36:46 +0000
Subject: [PATCH] qvm-run: wait for X11 in --dispvm --gui case

'qvm-run --dispvm' cannot easily make a separate qubes.WaitForSession
call. Instead, if --gui is active, pass the new WaitForSession argument
to qubes.VMShell, which will do the equivalent.

The unit tests have been copied (in slightly adapted form) from commit
https://github.com/marmarek/qubes-core-admin-client/commit/a620f02e2af8845805702f734188c7dd4027f734

Fixes QubesOS/qubes-issues#3012
Closes QubesOS/qubes-core-admin-client#49
---
 qubesadmin/tests/tools/qvm_run.py | 51 +++++++++++++++++++++++++++++++
 qubesadmin/tools/qvm_run.py       |  5 ++-
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/qubesadmin/tests/tools/qvm_run.py b/qubesadmin/tests/tools/qvm_run.py
index 0db6950..a052b95 100644
--- a/qubesadmin/tests/tools/qvm_run.py
+++ b/qubesadmin/tests/tools/qvm_run.py
@@ -428,3 +428,54 @@ class TC_00_qvm_run(qubesadmin.tests.QubesTestCase):
         self.assertEqual(ret, 0)
         self.assertEqual(self.app.service_calls, [])
         self.assertAllCalled()
+
+    def test_014_dispvm_local_gui(self):
+        self.app.qubesd_connection_type = 'socket'
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.CreateDisposable', None, None)] = \
+            b'0\0disp123'
+        self.app.expected_calls[('disp123', 'admin.vm.Kill', None, None)] = \
+            b'0\0'
+        self.app.expected_calls[
+            ('disp123', 'admin.vm.property.Get', 'qrexec_timeout', None)] = \
+            b'0\0default=yes type=int 30'
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--dispvm', '--', 'test.command'], app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('disp123', 'qubes.VMShell+WaitForSession', {
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+                'connect_timeout': 30,
+            }),
+            ('disp123', 'qubes.VMShell+WaitForSession',
+            b'test.command; exit\n'),
+        ])
+        self.assertAllCalled()
+
+    def test_015_dispvm_local_no_gui(self):
+        self.app.qubesd_connection_type = 'socket'
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.CreateDisposable', None, None)] = \
+            b'0\0disp123'
+        self.app.expected_calls[('disp123', 'admin.vm.Kill', None, None)] = \
+            b'0\0'
+        self.app.expected_calls[
+            ('disp123', 'admin.vm.property.Get', 'qrexec_timeout', None)] = \
+            b'0\0default=yes type=int 30'
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--dispvm', '--no-gui', 'test.command'], app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('disp123', 'qubes.VMShell', {
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+                'connect_timeout': 30,
+            }),
+            ('disp123', 'qubes.VMShell', b'test.command; exit\n'),
+        ])
+        self.assertAllCalled()
diff --git a/qubesadmin/tools/qvm_run.py b/qubesadmin/tools/qvm_run.py
index c3c1c7b..f10b6c2 100644
--- a/qubesadmin/tools/qvm_run.py
+++ b/qubesadmin/tools/qvm_run.py
@@ -209,7 +209,10 @@ def main(args=None, app=None):
                         localcmd=args.localcmd,
                         **run_kwargs)
                 else:
-                    proc = vm.run_service('qubes.VMShell',
+                    service = 'qubes.VMShell'
+                    if args.gui and args.dispvm:
+                        service += '+WaitForSession'
+                    proc = vm.run_service(service,
                         user=args.user,
                         localcmd=args.localcmd,
                         **run_kwargs)