Move qube actions to utils for reusability

Author: ben-grande

qubesadmin/tests/tools/qvm_shutdown.py

@@ -145,9 +145,7 @@ def test_015_wait_all_kill_timeout(self):
             b'sys-net class=AppVM state=Running\n' \
             b'some-vm class=AppVM state=Running\n' \
             b'other-vm class=AppVM state=Running\n'
-        with self.assertRaisesRegex(SystemExit, '2'):
-            qubesadmin.tools.qvm_shutdown.main(
-                ['--wait', '--all'], app=self.app)
+        qubesadmin.tools.qvm_shutdown.main(['--wait', '--all'], app=self.app)
         self.assertAllCalled()
 
     def test_016_all_exclude_noforce(self):
@@ -278,10 +276,10 @@ def test_011_wait_retry_once_fail(self):
             b'2\x00QubesVMInUseError\x00\x00Denied as qube is in use\x00',
             b'2\x00QubesVMInUseError\x00\x00Denied as qube is in use\x00',
         ]
-        with self.assertRaises(SystemExit) as cm:
+        with self.assertRaises(SystemExit) as cmg:
             qubesadmin.tools.qvm_shutdown.main(
                 ['--wait', 'some-vm', 'other-vm', 'another-vm'], app=self.app)
-        self.assertEqual(cm.exception.code, 3)
+        self.assertEqual(cmg.exception.code, 3)
         self.assertAllCalled()
 
     def test_011_wait_retry_as_many_as_in_use(self):
@@ -320,9 +318,10 @@ def test_013_wait_all_shutdown_fail(self):
         self.app.expected_calls[
             ('some-vm', 'admin.vm.Shutdown', 'wait', None)] = \
             b'2\x00QubesException\x00\x00Shutdown refused\x00'
-        with self.assertRaises(SystemExit):
+        with self.assertRaises(SystemExit) as cmg:
             qubesadmin.tools.qvm_shutdown.main(
                 ['--wait', 'some-vm'], app=self.app)
+        self.assertEqual(cmg.exception.code, 1)
         self.assertAllCalled()
 
     def test_014_timeout_deprecated_means_wait(self):
@@ -338,9 +337,10 @@ def test_014_timeout_deprecated_means_wait(self):
         self.app.expected_calls[
             ('some-vm', 'admin.vm.Kill', None, None)] = \
             b'2\x00QubesException\x00\x00Kill failed\x00'
-        with self.assertRaises(SystemExit):
+        with self.assertRaises(SystemExit) as cmg:
             qubesadmin.tools.qvm_shutdown.main(
                 ['--timeout=3', 'some-vm'], app=self.app)
+        self.assertEqual(cmg.exception.code, 1)
         self.assertAllCalled()
 
     def test_016_wait_kill_exception(self):
@@ -357,7 +357,8 @@ def test_016_wait_kill_exception(self):
         self.app.expected_calls[
             ('some-vm', 'admin.vm.Kill', None, None)] = \
             b'2\x00QubesException\x00\x00Kill failed\x00'
-        with self.assertRaises(SystemExit):
+        with self.assertRaises(SystemExit) as cmg:
             qubesadmin.tools.qvm_shutdown.main(
                 ['--wait', 'some-vm'], app=self.app)
+        self.assertEqual(cmg.exception.code, 1)
         self.assertAllCalled()

qubesadmin/tools/qvm_kill.py

@@ -37,15 +37,12 @@
 async def run_async(args=None, app=None):
     # pylint: disable=missing-docstring
     args = parser.parse_args(args, app=app)
-    remnants = await qubesadmin.tools.qvm_shutdown.kill(domains=args.domains)
-    if not remnants:
+    failed = await qubesadmin.utils.kill(domains=args.domains)
+    if not failed:
         return 0
-    parser.error_runtime(
-        "Failed to kill: {}".format(
-            ", ".join(qube.name for qube in remnants)
-        ),
-        len(remnants)
-    )
+    for qube, exc in failed.items():
+        parser.print_error("Failed to kill: {}: {}".format(qube, exc))
+    raise SystemExit(len(failed))
 
 
 def main(args=None, app=None):

qubesadmin/tools/qvm_pause.py

@@ -20,6 +20,7 @@
 
 '''qvm-pause - Pause a domain'''
 
+import asyncio
 import sys
 import qubesadmin
 
@@ -45,15 +46,16 @@ def main(args=None, app=None):
 
     args = parser.parse_args(args, app=app)
     exit_code = 0
-    for domain in args.domains:
-        try:
-            if args.suspend:
-                domain.suspend()
-            else:
-                domain.pause()
-        except (IOError, OSError, qubesadmin.exc.QubesException) as e:
-            exit_code = 1
-            parser.print_error(str(e))
+    if args.suspend:
+        action = "suspend"
+        failed = asyncio.run(qubesadmin.utils.suspend(domains=args.domains))
+    else:
+        action = "pause"
+        failed = asyncio.run(qubesadmin.utils.pause(domains=args.domains))
+    if failed:
+        exit_code = 1
+        for qube, exc in failed.items():
+            parser.print_error("Failed to {}: {}: {}".format(action, qube, exc))
 
     return exit_code
 

qubesadmin/tools/qvm_shutdown.py

@@ -67,67 +67,27 @@
 )
 
 
-async def shutdown(args, domains: list[qubesadmin.vm.QubesVM]):
-    """
-    Asynchronously shutdown qubes and return qubes that failed to shutdown
-    because and the client can't handle, as well as qubes that were in use
-    while --force was not provided, as well as timed out.
-    """
-    # pylint: disable=missing-docstring
-    unhandled, used, timedout = [], [], []
-    tasks = [
-        asyncio.to_thread(qube.shutdown, force=args.force, wait=args.wait)
-        for qube in domains
-    ]
-    results = await asyncio.gather(*tasks, return_exceptions=True)
-    for qube, res in zip(domains, results):
-        if not isinstance(res, BaseException):
-            qube.log.info("Shutdown succeeded")
-            continue
-        try:
-            raise res
-        except qubesadmin.exc.QubesVMNotStartedError:
-            pass
-        except qubesadmin.exc.QubesVMInUseError as e:
-            if args.wait:
-                qube.log.error("Shutdown error: {}".format(e))
-            else:
-                qube.log.error("Shutdown error: (try --force): {}".format(e))
-            used.append(qube)
-        except qubesadmin.exc.QubesVMShutdownTimeoutError as e:
-            if args.wait:
-                qube.log.error("Shutdown error: {}".format(e))
-            else:
-                qube.log.error("Shutdown error: (try qvm-kill): {}".format(e))
-            timedout.append(qube)
-        except qubesadmin.exc.QubesException as e:
-            qube.log.error("Shutdown error: {}".format(e))
-            unhandled.append(qube)
+async def shutdown(domains, **shutdown_kwargs):
+    # pylint: disable=missing-function-docstring
+    failed = await qubesadmin.utils.shutdown(domains=domains, **shutdown_kwargs)
+    used = {
+        qube: exc
+        for qube, exc in failed.items()
+        if isinstance(exc, qubesadmin.exc.QubesVMInUseError)
+    }
+    timedout = {
+        qube: exc
+        for qube, exc in failed.items()
+        if isinstance(exc, qubesadmin.exc.QubesVMShutdownTimeoutError)
+    }
+    unhandled = {
+        qube: exc
+        for qube, exc in failed.items()
+        if qube not in used and qube not in timedout
+    }
     return unhandled, used, timedout
 
 
-async def kill(domains: list[qubesadmin.vm.QubesVM]):
-    """
-    Asynchronously kill qubes and return qubes that failed to shutdown.
-    """
-    # pylint: disable=missing-docstring
-    unhandled = domains.copy()
-    tasks = [asyncio.to_thread(qube.kill) for qube in domains]
-    results = await asyncio.gather(*tasks, return_exceptions=True)
-    for qube, res in zip(domains, results):
-        if not isinstance(res, BaseException):
-            qube.log.info("Killing succeeded")
-            unhandled.remove(qube)
-            continue
-        try:
-            raise res
-        except qubesadmin.exc.QubesVMNotStartedError:
-            unhandled.remove(qube)
-        except qubesadmin.exc.QubesException as e:
-            qube.log.error("Kill error: {}".format(e))
-    return unhandled
-
-
 async def run_async(args=None, app=None):
     # pylint: disable=missing-docstring
     args = parser.parse_args(args, app=app)
@@ -140,10 +100,16 @@ async def run_async(args=None, app=None):
         )
         args.wait = True
     args.force = args.force or (args.all_domains and not args.exclude)
-
-    unhandled, used, timedout = await shutdown(args=args, domains=args.domains)
-    unhandled_retry = []
-    timedout_retry = []
+    shutdown_kwargs = {
+        "force": args.force,
+        "wait": args.wait,
+    }
+
+    unhandled, used, timedout = await shutdown(
+        domains=args.domains, **shutdown_kwargs
+    )
+    unhandled_retry = {}
+    timedout_retry = {}
     if used:
         old_failed = unhandled, used, timedout
         parser.print_error(
@@ -157,7 +123,7 @@ async def run_async(args=None, app=None):
                 ", ".join(qube.name for qube in used)
             )
         )
-        failed = await shutdown(args=args, domains=used)
+        failed = await shutdown(domains=used, **shutdown_kwargs)
         unhandled_retry, used, timedout_retry = failed
         if not failed:
             break
@@ -168,8 +134,8 @@ async def run_async(args=None, app=None):
                 break
         old_failed = failed
 
-    unhandled.extend(qube for qube in unhandled_retry if qube not in unhandled)
-    timedout.extend(qube for qube in timedout_retry if qube not in timedout)
+    unhandled.update(unhandled_retry)
+    timedout.update(timedout_retry)
 
     # Retry timed out only once, as it can take a long time, 60s by default.
     if timedout:
@@ -178,38 +144,28 @@ async def run_async(args=None, app=None):
                 ", ".join(qube.name for qube in timedout)
             )
         )
-        unhandled, used, timedout = await shutdown(args=args, domains=timedout)
+        unhandled, used, timedout = await shutdown(
+            domains=timedout, **shutdown_kwargs
+        )
 
     if timedout:
         parser.print_error(
             "Killing timed out qubes: {}".format(
                 ", ".join(qube.name for qube in timedout)
             )
         )
-        unhandled = await kill(domains=timedout)
+        timedout = await qubesadmin.utils.kill(domains=timedout)
 
-    if not unhandled and not used and not timedout:
-        return
-
-    if unhandled:
-        parser.print_error(
-            "Failed to shut down for unknown reason: {}".format(
-                ", ".join(qube.name for qube in unhandled)
-            )
-        )
-    if used:
-        parser.print_error(
-            "Failed to shut down because it's in use: {}".format(
-                ", ".join(qube.name for qube in used)
+    for item in [unhandled, used, timedout]:
+        for qube, exc in item.items():
+            parser.print_error(
+                "Failed to shut down: {}: {}".format(qube.name, str(exc))
             )
-        )
-    if timedout:
-        parser.print_error(
-            "Failed to shut down because of time out: {}".format(
-                ", ".join(qube.name for qube in timedout)
-            )
-        )
-    raise SystemExit(len(unhandled + used + timedout))
+
+    exit_code = len(unhandled) + len(used) + len(timedout)
+    if exit_code == 0:
+        return
+    raise SystemExit(exit_code)
 
 
 def main(args=None, app=None):