tools: add qvm-run --dispvm option

Add option to uniformly start new DispVM from either VM or Dom0. This
use DispVMWrapper, which translate it to either qrexec call to $dispvm,
or (in dom0) to appropriate Admin API call to create fresh DispVM
first.
This require abandoning registering --all and --exclude by
QubesArgumentParser, because we need to add --dispvm mutually exclusive
with those two. But actually handling those two options is still done by
QubesArgumentParser.

This also updates man page and tests.

Fixes QubesOS/qubes-issues#2974

doc/manpages/qvm-run.rst

@@ -6,7 +6,9 @@
 Synopsis
 --------
 
-:command:`qvm-run` [-h] [--verbose] [--quiet] [--all] [--exclude *EXCLUDE*] [--user *USER*] [--autostart] [--pass-io] [--localcmd *COMMAND*] [--gui] [--no-gui] [--colour-output *COLOR*] [--no-color-output] [--filter-escape-chars] [--no-filter-escape-chars] [*VMNAME*] *COMMAND*
+:command:`qvm-run` [options] *VMNAME* *COMMAND*
+:command:`qvm-run` [options] --all [--exclude *EXCLUDE*]  *COMMAND*
+:command:`qvm-run` [options] --dispvm [*BASE_APPVM*] *COMMAND*
 
 Options
 -------
@@ -32,6 +34,11 @@ Options
 
    Exclude the qube from :option:`--all`.
 
+.. option:: --dispvm [BASE_APPVM]
+
+   Run the command fresh DisposableVM created out of *BASE_APPVM*. This option
+   is mutually exclusive with *VMNAME*, --all and --exclude.
+
 .. option:: --user=USER, -u USER
 
    Run command in a qube as *USER*.

qubesadmin/tests/tools/qvm_run.py

@@ -66,12 +66,19 @@ class TC_00_qvm_run(qubesadmin.tests.QubesTestCase):
         self.app.expected_calls[
             ('dom0', 'admin.vm.List', None, None)] = \
             b'0\x00test-vm class=AppVM state=Running\n' \
-            b'test-vm2 class=AppVM state=Running\n'
+            b'test-vm2 class=AppVM state=Running\n' \
-        # self.app.expected_calls[
+            b'test-vm3 class=AppVM state=Halted\n'
-        #     ('test-vm', 'admin.vm.List', None, None)] = \
+        self.app.expected_calls[
-        #     b'0\x00test-vm class=AppVM state=Running\n'
+            ('test-vm', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm class=AppVM state=Running\n'
+        self.app.expected_calls[
+            ('test-vm2', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm2 class=AppVM state=Running\n'
+        self.app.expected_calls[
+            ('test-vm3', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm3 class=AppVM state=Halted\n'
         ret = qubesadmin.tools.qvm_run.main(
-            ['--no-gui', 'test-vm', 'test-vm2', 'command'],
+            ['--no-gui', '--all', 'command'],
             app=self.app)
         self.assertEqual(ret, 0)
         self.assertEqual(self.app.service_calls, [
@@ -304,3 +311,107 @@ class TC_00_qvm_run(qubesadmin.tests.QubesTestCase):
             ('test-vm', 'service.name', b''),
         ])
         self.assertAllCalled()
+
+    def test_008_dispvm_remote(self):
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--dispvm', '--service', 'test.service'], app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('$dispvm', 'test.service', {
+                'filter_esc': self.default_filter_esc(),
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+            }),
+            ('$dispvm', 'test.service', b''),
+        ])
+        self.assertAllCalled()
+
+    def test_009_dispvm_remote_specific(self):
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--dispvm=test-vm', '--service', 'test.service'], app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('$dispvm:test-vm', 'test.service', {
+                'filter_esc': self.default_filter_esc(),
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+            }),
+            ('$dispvm:test-vm', 'test.service', b''),
+        ])
+        self.assertAllCalled()
+
+    def test_010_dispvm_local(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'
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--dispvm', '--service', 'test.service'], app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('disp123', 'test.service', {
+                'filter_esc': self.default_filter_esc(),
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+            }),
+            ('disp123', 'test.service', b''),
+        ])
+        self.assertAllCalled()
+
+    def test_011_dispvm_local_specific(self):
+        self.app.qubesd_connection_type = 'socket'
+        self.app.expected_calls[
+            ('test-vm', 'admin.vm.CreateDisposable', None, None)] = \
+            b'0\0disp123'
+        self.app.expected_calls[('disp123', 'admin.vm.Kill', None, None)] = \
+            b'0\0'
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--dispvm=test-vm', '--service', 'test.service'], app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('disp123', 'test.service', {
+                'filter_esc': self.default_filter_esc(),
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+            }),
+            ('disp123', 'test.service', b''),
+        ])
+        self.assertAllCalled()
+
+    def test_012_exclude(self):
+        self.app.expected_calls[
+            ('dom0', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm class=AppVM state=Running\n' \
+            b'test-vm2 class=AppVM state=Running\n' \
+            b'test-vm3 class=AppVM state=Halted\n'
+        self.app.expected_calls[
+            ('test-vm', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm class=AppVM state=Running\n'
+        self.app.expected_calls[
+            ('test-vm3', 'admin.vm.List', None, None)] = \
+            b'0\x00test-vm3 class=AppVM state=Halted\n'
+        ret = qubesadmin.tools.qvm_run.main(
+            ['--no-gui', '--all', '--exclude', 'test-vm2', 'command'],
+            app=self.app)
+        self.assertEqual(ret, 0)
+        self.assertEqual(self.app.service_calls, [
+            ('test-vm', 'qubes.VMShell', {
+                'filter_esc': self.default_filter_esc(),
+                'localcmd': None,
+                'stdout': subprocess.DEVNULL,
+                'stderr': subprocess.DEVNULL,
+                'user': None,
+            }),
+            ('test-vm', 'qubes.VMShell', b'command; exit\n'),
+        ])
+        self.assertAllCalled()

qubesadmin/tools/qvm_run.py

@@ -30,7 +30,7 @@ import multiprocessing
 import qubesadmin.tools
 import qubesadmin.exc
 
-parser = qubesadmin.tools.QubesArgumentParser(vmname_nargs='+')
+parser = qubesadmin.tools.QubesArgumentParser()
 
 parser.add_argument('--user', '-u', metavar='USER',
     help='run command in a qube as USER (available only from dom0)')
@@ -88,6 +88,24 @@ parser.add_argument('--service',
     action='store_true', dest='service',
     help='run a qrexec service (named by COMMAND) instead of shell command')
 
+target_parser = parser.add_mutually_exclusive_group()
+
+target_parser.add_argument('--dispvm', action='store', nargs='?',
+    const=True, metavar='BASE_APPVM',
+    help='start a service in new Disposable VM; '
+         'optionally specify base AppVM for DispVM')
+target_parser.add_argument('VMNAME',
+    nargs='?',
+    action=qubesadmin.tools.VmNameAction)
+
+# add those manually instead of vmname_args, because of mutually exclusive
+# group with --dispvm; parsing is still handled by QubesArgumentParser
+target_parser.add_argument('--all', action='store_true', dest='all_domains',
+    help='run command on all running qubes')
+
+parser.add_argument('--exclude', action='append', default=[],
+    help='exclude the qube from --all')
+
 parser.add_argument('cmd', metavar='COMMAND',
     help='command to run')
 
@@ -130,7 +148,10 @@ def main(args=None, app=None):
         run_kwargs['stderr'] = None
 
     if isinstance(args.app, qubesadmin.app.QubesLocal) and \
-            not args.passio and not args.localcmd and args.service:
+            not args.passio and \
+            not args.localcmd and \
+            args.service and \
+            not args.dispvm:
         # wait=False works only in dom0; but it's still useful, to save on
         # simultaneous vchan connections
         run_kwargs['wait'] = False
@@ -139,6 +160,18 @@ def main(args=None, app=None):
     if args.passio:
         verbose -= 1
 
+    # --all and --exclude are handled by QubesArgumentParser
+    domains = args.domains
+    dispvm = None
+    if args.dispvm:
+        if args.exclude:
+            parser.error('Cannot use --exclude with --dispvm')
+        dispvm = qubesadmin.vm.DispVM.from_appvm(args.app,
+            None if args.dispvm is True else args.dispvm)
+        domains = [dispvm]
+    elif args.all_domains:
+        # --all consider only running VMs
+        domains = [vm for vm in domains if vm.is_running()]
     if args.color_output:
         sys.stdout.write('\033[0;{}m'.format(args.color_output))
         sys.stdout.flush()
@@ -148,7 +181,7 @@ def main(args=None, app=None):
     copy_proc = None
     try:
         procs = []
-        for vm in args.domains:
+        for vm in domains:
             if not args.autostart and not vm.is_running():
                 continue
             try:
@@ -160,7 +193,7 @@ def main(args=None, app=None):
                     else:
                         print('Running \'{}\' on {}'.format(args.cmd, vm.name),
                             file=sys.stderr)
-                if args.gui:
+                if args.gui and not args.dispvm:
                     wait_session = vm.run_service('qubes.WaitForSession',
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                     wait_session.communicate(vm.default_user.encode())
@@ -194,6 +227,8 @@ def main(args=None, app=None):
         for proc in procs:
             retcode = max(retcode, proc.wait())
     finally:
+        if dispvm:
+            dispvm.cleanup()
         if args.color_output:
             sys.stdout.write('\033[0m')
             sys.stdout.flush()