diff --git a/qubesmanager/qube_manager.py b/qubesmanager/qube_manager.py index 7da14e9e..d0157da1 100644 --- a/qubesmanager/qube_manager.py +++ b/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,6 +1397,10 @@ 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) @@ -1390,7 +1408,7 @@ def table_selection_changed(self): 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() diff --git a/qubesmanager/tests/test_qube_manager.py b/qubesmanager/tests/test_qube_manager.py index 09632868..e174e5fc 100644 --- a/qubesmanager/tests/test_qube_manager.py +++ b/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')