From d7c3ee93149201f06dbaf97ceeeb871f9eced98e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Pierret=20=28fepitre=29?=
 <frederic.pierret@qubes-os.org>
Date: Tue, 20 Aug 2024 15:31:37 +0200
Subject: [PATCH] Improve interrupt handling and send signal to current loop

Closes QubesOS/qubes-issues#9410
---
 qubesbuilder/cli/cli_base.py    | 43 +++++++++++++++++++++++----------
 qubesbuilder/executors/qubes.py |  2 +-
 tests/test_cli.py               | 34 +++++++++++++++++++++++++-
 3 files changed, 64 insertions(+), 15 deletions(-)

diff --git a/qubesbuilder/cli/cli_base.py b/qubesbuilder/cli/cli_base.py
index cfed5f05..0b921379 100644
--- a/qubesbuilder/cli/cli_base.py
+++ b/qubesbuilder/cli/cli_base.py
@@ -20,7 +20,8 @@
 """
 QubesBuilder command-line interface - base module.
 """
-
+import asyncio
+import signal
 import sys
 import traceback
 from typing import Callable, List
@@ -61,21 +62,37 @@ def __init__(self, *args, **kwargs):
         self.list_commands = self.list_commands_for_help  # type: ignore
 
     def __call__(self, *args, **kwargs):
+        loop = asyncio.get_event_loop()
+
+        def _handle_interrupt():
+            tasks = asyncio.all_tasks(loop)
+            for task in tasks:
+                task.cancel()
+
+        loop.add_signal_handler(signal.SIGINT, _handle_interrupt)
+
+        rc = 1
         try:
-            return self.main(*args, **kwargs)
+            rv = self.main(*args, standalone_mode=False, **kwargs)
+            if isinstance(rv, list) and set(rv) == {None}:
+                rc = 0
         except Exception as exc:
-            QubesBuilderLogger.error(f"An error occurred: {str(exc)}")
-            if self.debug is True:
-                formatted_traceback = "".join(traceback.format_exception(exc))
-                QubesBuilderLogger.error(
-                    "\n" + formatted_traceback.rstrip("\n")
-                )
-
-            if isinstance(exc, click.ClickException):
-                # pylint: disable=no-member
-                sys.exit(exc.exit_code)
+            if isinstance(exc, click.Abort) or "Signals.SIGINT" in str(exc):
+                QubesBuilderLogger.warning(f"Interrupting...")
             else:
-                sys.exit(1)
+                QubesBuilderLogger.error(f"An error occurred: {str(exc)}")
+                if self.debug is True:
+                    formatted_traceback = "".join(
+                        traceback.format_exception(exc)
+                    )
+                    QubesBuilderLogger.error(
+                        "\n" + formatted_traceback.rstrip("\n")
+                    )
+                if isinstance(exc, click.ClickException):
+                    # pylint: disable=no-member
+                    rc = exc.exit_code
+        finally:
+            sys.exit(rc)
 
     def get_command(self, ctx, cmd_name):
         rv = click.Group.get_command(self, ctx, cmd_name)
diff --git a/qubesbuilder/executors/qubes.py b/qubesbuilder/executors/qubes.py
index 8dd0d652..408577a5 100644
--- a/qubesbuilder/executors/qubes.py
+++ b/qubesbuilder/executors/qubes.py
@@ -371,7 +371,7 @@ def run(  # type: ignore
                         )
                         continue
                     raise e
-        except ExecutorError as e:
+        except (subprocess.CalledProcessError, ExecutorError) as e:
             if dispvm and self._clean_on_error:
                 self.cleanup(dispvm)
             raise e
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 84bec3d3..3b4fbda2 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -43,7 +43,7 @@ def qb_call(builder_conf, artifacts_dir, *args, **kwargs):
         f"artifacts-dir={artifacts_dir}",
         *args,
     ]
-    subprocess.check_call(cmd, **kwargs)
+    return subprocess.check_call(cmd, **kwargs)
 
 
 def qb_call_output(builder_conf, artifacts_dir, *args, **kwargs):
@@ -2567,3 +2567,35 @@ def test_installer_init_cache(artifacts_dir):
     rpms = list(templates_cache.glob("*.rpm"))
     assert rpms
     assert rpms[0].name.startswith("qubes-template-debian-12-minimal-4.2.0")
+
+
+def test_existent_command(artifacts_dir):
+    result = qb_call(
+        DEFAULT_BUILDER_CONF,
+        artifacts_dir,
+        "config",
+        "get-components",
+        "get-templates",
+    )
+    assert result == 0
+
+
+def test_non_existent_command(artifacts_dir):
+    with pytest.raises(subprocess.CalledProcessError):
+        result = qb_call(
+            DEFAULT_BUILDER_CONF, artifacts_dir, "non-existent-command"
+        )
+        assert result == 2
+
+
+def test_non_existent_component(artifacts_dir):
+    with pytest.raises(subprocess.CalledProcessError):
+        result = qb_call(
+            DEFAULT_BUILDER_CONF,
+            artifacts_dir,
+            "-c",
+            "non-existent-component",
+            "package",
+            "all",
+        )
+        assert result == 2