qui-domains: adjust shutdown failure actions

Author: Jayant-kernel

qui/tray/domains.py

@@ -192,21 +192,36 @@ async def perform_action(self):
 class ShutdownItem(VMActionMenuItem):
     """Shutdown menu Item. When activated shutdowns the domain."""
 
-    def __init__(self, vm, icon_cache, force=False):
+    def __init__(
+        self,
+        vm,
+        icon_cache,
+        force=False,
+        follow_shift=True,
+        label=None,
+        force_label=None,
+        icon_name="shutdown",
+    ):
         if force:
             super().__init__(
                 vm,
-                label=_("Force shutdown"),
+                label=force_label or _("Force shutdown"),
                 icon_cache=icon_cache,
-                icon_name="shutdown",
+                icon_name=icon_name,
             )
         else:
             super().__init__(
-                vm, label=_("Shutdown"), icon_cache=icon_cache, icon_name="shutdown"
+                vm,
+                label=label or _("Shutdown"),
+                icon_cache=icon_cache,
+                icon_name=icon_name,
             )
         self.force = force
+        self.follow_shift = follow_shift
 
     def set_force(self, force):
+        if not self.follow_shift:
+            return
         self.force = force
         if self.force:
             self.label.set_text(_("Force shutdown"))
@@ -226,89 +241,97 @@ async def perform_action(self):
                     ).format(self.vm.name, str(ex)),
                 )
                 return
-            if isinstance(ex, exc.QubesVMInUseError):
-                title = _("Qube {0} is in use").format(self.vm.name)
-                markup = _(
-                    "The qube <b>{0}</b> could not be shut down because it "
-                    "is in use by connected qubes.\n\n"
-                    "<b>Warning:</b> force shutdown may cause unexpected "
-                    "issues in connected qubes."
-                ).format(self.vm.name)
-                button_label = _("Force shutdown")
-                action = "force"
-            elif isinstance(ex, exc.QubesVMShutdownTimeoutError):
-                title = _("Qube {0} shutdown timed out").format(self.vm.name)
-                markup = _(
-                    "The qube <b>{0}</b> did not shut down within the "
-                    "expected time.\n\nYou can retry the shutdown or, "
-                    "if the problem persists, kill the qube."
-                ).format(self.vm.name)
-                button_label = None  # timeout uses custom buttons below
-                action = "timeout"
-            else:
-                title = _("Error shutting down qube {0}").format(self.vm.name)
-                markup = _(
-                    "The qube <b>{0}</b> could not be shut down. "
-                    "The following error occurred:\n"
-                    "<tt>{1}</tt>\n\n"
-                    "Would you like to kill the qube?"
-                ).format(self.vm.name, str(ex))
-                button_label = _("Kill")
-                action = "kill"
-
-            dialog = Gtk.MessageDialog(
-                None, 0, Gtk.MessageType.ERROR, Gtk.ButtonsType.NONE
-            )
-            dialog.set_title(title)
-            dialog.set_markup(markup)
-            dialog.add_button(_("Cancel"), Gtk.ResponseType.CANCEL)
-            if action == "timeout":
-                dialog.add_button(_("Retry shutdown"), Gtk.ResponseType.OK)
-                dialog.add_button(_("Kill"), Gtk.ResponseType.YES)
-            else:
-                dialog.add_button(button_label, Gtk.ResponseType.OK)
-            dialog.connect("response", self.react_to_question, action)
-            GLib.idle_add(dialog.show)
+            self.show_shutdown_dialog(ex)
+
+    def get_shutdown_question(self, ex):
+        if isinstance(ex, exc.QubesVMInUseError):
+            title = _("Qube {0} is in use").format(self.vm.name)
+            markup = _(
+                "The qube <b>{0}</b> could not be shut down because it "
+                "is in use by connected qubes.\n\n"
+                "<b>Warning:</b> force shutdown may cause unexpected "
+                "issues in connected qubes."
+            ).format(self.vm.name)
+            action = "force"
+        elif isinstance(ex, exc.QubesVMShutdownTimeoutError):
+            title = _("Qube {0} shutdown timed out").format(self.vm.name)
+            markup = _(
+                "The qube <b>{0}</b> did not shut down within the "
+                "expected time.\n\nYou can retry the shutdown or, "
+                "if the problem persists, kill the qube."
+            ).format(self.vm.name)
+            action = "timeout"
+        else:
+            title = _("Error shutting down qube {0}").format(self.vm.name)
+            markup = _(
+                "The qube <b>{0}</b> could not be shut down. "
+                "The following error occurred:\n"
+                "<tt>{1}</tt>\n\n"
+                "Would you like to kill the qube?"
+            ).format(self.vm.name, str(ex))
+            action = "kill"
+        return title, markup, action
+
+    def show_shutdown_dialog(self, ex):
+        title, markup, action = self.get_shutdown_question(ex)
+        dialog = Gtk.MessageDialog(None, 0, Gtk.MessageType.ERROR, Gtk.ButtonsType.NONE)
+        dialog.set_title(title)
+        dialog.set_markup(markup)
+        dialog.add_button(_("Cancel"), Gtk.ResponseType.CANCEL)
+        if action == "force":
+            dialog.add_button(_("Force shutdown"), Gtk.ResponseType.OK)
+        elif action == "timeout":
+            dialog.add_button(_("Retry shutdown"), Gtk.ResponseType.OK)
+            dialog.add_button(_("Kill"), Gtk.ResponseType.YES)
+        elif action == "kill":
+            dialog.add_button(_("Kill"), Gtk.ResponseType.OK)
+        dialog.connect("response", self.react_to_question, action)
+        GLib.idle_add(dialog.show)
 
     def react_to_question(self, widget, response, action):
         if response not in (Gtk.ResponseType.OK, Gtk.ResponseType.YES):
             widget.destroy()
             return
         try:
-            if action == "force":
-                self.vm.shutdown(force=True)
-            elif action == "timeout":
-                if response == Gtk.ResponseType.YES:
-                    self.vm.kill()
-                elif response == Gtk.ResponseType.OK:
-                    self.vm.shutdown(force=False)
-            elif action == "kill" and response == Gtk.ResponseType.OK:
-                self.vm.kill()
+            self.do_shutdown_dialog_action(response, action)
         except exc.QubesException as ex:
-            show_error(
-                _("Error shutting down qube"),
-                _(
-                    "The following error occurred while attempting to "
-                    "shut down qube {0}:\n{1}"
-                ).format(self.vm.name, str(ex)),
-            )
+            if action == "timeout" and response == Gtk.ResponseType.OK:
+                self.show_shutdown_dialog(ex)
+            else:
+                show_error(
+                    _("Error shutting down qube"),
+                    _(
+                        "The following error occurred while attempting to "
+                        "shut down qube {0}:\n{1}"
+                    ).format(self.vm.name, str(ex)),
+                )
         widget.destroy()
 
+    def do_shutdown_dialog_action(self, response, action):
+        if action == "force":
+            self.vm.shutdown(force=True)
+        elif action == "timeout":
+            if response == Gtk.ResponseType.YES:
+                self.vm.kill()
+            elif response == Gtk.ResponseType.OK:
+                self.vm.shutdown(force=False)
+        elif action == "kill" and response == Gtk.ResponseType.OK:
+            self.vm.kill()
+
 
-class RestartItem(VMActionMenuItem):
+class RestartItem(ShutdownItem):
     """Restart menu Item. When activated shutdowns the domain and
     then starts it again."""
 
     def __init__(self, vm, icon_cache, force=False):
-        if force:
-            super().__init__(
-                vm, label=_("Force restart"), icon_cache=icon_cache, icon_name="restart"
-            )
-        else:
-            super().__init__(
-                vm, label=_("Restart"), icon_cache=icon_cache, icon_name="restart"
-            )
-        self.force = force
+        super().__init__(
+            vm,
+            icon_cache,
+            force=force,
+            label=_("Restart"),
+            force_label=_("Force restart"),
+            icon_name="restart",
+        )
         self.give_up = False
 
     def set_force(self, force):
@@ -332,19 +355,7 @@ async def perform_action(self, *_args, **_kwargs):
                     ).format(self.vm.name, str(ex)),
                 )
                 return
-            dialog = Gtk.MessageDialog(
-                None, 0, Gtk.MessageType.ERROR, Gtk.ButtonsType.OK_CANCEL
-            )
-            dialog.set_title("Error restarting qube")
-            dialog.set_markup(
-                f"The qube {self.vm.name} couldn't be shut down "
-                "normally. The following error occurred: \n"
-                f"<tt>{str(ex)}</tt>\n\n"
-                "Do you want to force shutdown? \n\n<b>Warning:</b> "
-                "this may cause unexpected issues in connected qubes."
-            )
-            dialog.connect("response", self.react_to_question)
-            GLib.idle_add(dialog.show)
+            self.show_shutdown_dialog(ex)
 
         try:
             while self.vm.is_running():
@@ -366,11 +377,17 @@ async def perform_action(self, *_args, **_kwargs):
                 ).format(self.vm.name, str(ex)),
             )
 
-    def react_to_question(self, widget, response):
-        if response == Gtk.ResponseType.OK:
-            try:
-                self.vm.shutdown(force=True)
-            except exc.QubesException as ex:
+    def react_to_question(self, widget, response, action):
+        if response not in (Gtk.ResponseType.OK, Gtk.ResponseType.YES):
+            self.give_up = True
+            widget.destroy()
+            return
+        try:
+            self.do_shutdown_dialog_action(response, action)
+        except exc.QubesException as ex:
+            if action == "timeout" and response == Gtk.ResponseType.OK:
+                self.show_shutdown_dialog(ex)
+            else:
                 show_error(
                     _("Error shutting down qube"),
                     _(
@@ -379,8 +396,6 @@ def react_to_question(self, widget, response):
                     ).format(self.vm.name, str(ex)),
                 )
                 self.give_up = True
-        else:
-            self.give_up = True
         widget.destroy()
 
 
@@ -547,7 +562,7 @@ def __init__(self, is_preload=False):
 class StartedMenu(Gtk.Menu):
     """The sub-menu for a started domain"""
 
-    def __init__(self, vm, app, icon_cache):
+    def __init__(self, vm, app, icon_cache, shutdown_failed=False):
         super().__init__()
         self.vm = vm
         self.app = app
@@ -561,7 +576,17 @@ def __init__(self, vm, app, icon_cache):
 
         self.add(PreferencesItem(self.vm, icon_cache))
         self.add(PauseItem(self.vm, icon_cache))
-        self.add(ShutdownItem(self.vm, icon_cache, force=app.shift_pressed))
+        self.add(
+            ShutdownItem(
+                self.vm,
+                icon_cache,
+                force=app.shift_pressed and not shutdown_failed,
+                follow_shift=not shutdown_failed,
+            )
+        )
+        if shutdown_failed:
+            self.add(ShutdownItem(self.vm, icon_cache, force=True, follow_shift=False))
+            self.add(KillItem(self.vm, icon_cache))
         if self.vm.klass != "DispVM" or not self.vm.auto_cleanup:
             self.add(RestartItem(self.vm, icon_cache, force=app.shift_pressed))
 
@@ -709,6 +734,7 @@ def __init__(self, vm, app, icon_cache, state=None):
         self.vm = vm
         self.app = app
         self.icon_cache = icon_cache
+        self.shutdown_failed = False
         self.decorator = qui.decorators.DomainDecorator(vm)
 
         # Main horizontal box
@@ -766,7 +792,12 @@ def _set_submenu(self, state):
                 is_preload=is_preload,
             )
         elif state == "Running":
-            submenu = StartedMenu(self.vm, self.app, self.icon_cache)
+            submenu = StartedMenu(
+                self.vm,
+                self.app,
+                self.icon_cache,
+                shutdown_failed=self.shutdown_failed,
+            )
         elif state == "Paused":
             submenu = PausedMenu(self.vm, self.icon_cache)
         else:
@@ -1182,6 +1213,11 @@ def update_domain_item(self, vm, event, **kwargs):
                 # it's a fragile DispVM
                 state = "Transient"
 
+        if event == "domain-shutdown-failed":
+            item.shutdown_failed = True
+        elif event in ("domain-start", "domain-pre-start", "domain-shutdown"):
+            item.shutdown_failed = False
+
         item.update_state(state)
 
         if event == "property-reset:is_preload":