Fix changing templates for disposable qubes

qubesmanager/qube_manager.py

@@ -1141,13 +1141,27 @@ def fill_cache(self):
 
         progress.setValue(row_no)
 
-    def init_template_menu(self):
+    def init_template_menu(self, selected_vms=None):
         self.template_menu.clear()
+
+        dispvm_selection = (
+            bool(selected_vms)
+            and all(vm.klass == 'DispVM' for vm in selected_vms)
+        )
+
         for vm in self.qubes_app.domains:
-            if vm.klass == 'TemplateVM':
+            if (
+                dispvm_selection
+                and getattr(vm, 'template_for_dispvms', False)
+            ) or (
+                not dispvm_selection
+                and vm.klass == 'TemplateVM'
+            ):
                 action = self.template_menu.addAction(vm.name)
                 action.setData(vm.name)
-                action.triggered.connect(partial(self.change_template, vm.name))
+                action.triggered.connect(
+                    partial(self.change_template, vm.name)
+                )
 
     def _get_default_netvm(self):
         for vm in self.qubes_app.domains:
@@ -1383,14 +1397,18 @@ def get_selected_vms(self):
         return vms
 
     def table_selection_changed(self):
+        selected_vms = self.get_selected_vms()
+
+        self.init_template_menu(selected_vms)
+
         # Since selection could have multiple domains
         # enable all first and then filter them
         self.template_menu.setEnabled(True)
         self.network_menu.setEnabled(True)
         for action in self.toolbar.actions() + self.context_menu.actions():
             action.setEnabled(True)
 
-        for vm in self.get_selected_vms():
+        for vm in selected_vms:
             #  TODO: add boot from device to menu and add windows tools there
             # Update available actions:
             if vm.state['power'] in \
@@ -1452,7 +1470,7 @@ def table_selection_changed(self):
                 self.action_appmenus.setEnabled(False)
                 if vm.auto_cleanup:
                     self.action_restartvm.setEnabled(False)
-                self.template_menu.setEnabled(False)
+
             elif vm.klass == 'TemplateVM':
                 self.template_menu.setEnabled(False)
                 self.network_menu.setEnabled(False)
@@ -1474,6 +1492,12 @@ def table_selection_changed(self):
                 self.action_updatevm.setEnabled(False)
                 self.action_run_command_in_vm.setEnabled(False)
 
+        if (
+            any(vm.klass == 'DispVM' for vm in selected_vms)
+            and not all(vm.klass == 'DispVM' for vm in selected_vms)
+        ):
+            self.template_menu.setEnabled(False)
+
         self.update_template_menu()
         self.update_network_menu()
 

qubesmanager/tests/test_qube_manager.py

@@ -985,6 +985,45 @@ def test_313_template_menu_multiple(mock_question, qubes_manager):
     for call in calls:
         assert call in qubes_manager.qubes_app.actual_calls
 
+@pytest.mark.asyncio(loop_scope="module")
+@mock.patch('PyQt6.QtWidgets.QMessageBox.question')
+async def test_312_template_menu_dispvm(mock_question, qubes_manager):
+    mock_question.return_value = QMessageBox.StandardButton.Yes
+
+    _select_vm(qubes_manager, 'test-disp')
+
+    assert qubes_manager.template_menu.isEnabled()
+
+    expected_templates = {
+        str(vm) for vm in qubes_manager.qubes_app.domains
+        if getattr(vm, 'template_for_dispvms', False)
+    }
+
+    current_templates = {
+        action.text() for action in qubes_manager.template_menu.actions()
+    }
+
+    assert current_templates == expected_templates
+
+    change_call = (
+        'test-disp',
+        'admin.vm.property.Set',
+        'template',
+        b'test-alt-dvm'
+    )
+
+    qubes_manager.qubes_app.expected_calls[change_call] = b'0\x00'
+
+    action = next(
+        action for action in qubes_manager.template_menu.actions()
+        if action.text() == 'test-alt-dvm'
+    )
+
+    action.trigger()
+
+    await asyncio.sleep(0)
+
+    assert change_call in qubes_manager.qubes_app.actual_calls
 
 @mock.patch('PyQt6.QtWidgets.QMessageBox.information')
 @mock.patch('PyQt6.QtWidgets.QMessageBox.warning')