Island: Create callable reset_agent_configuration class

monkey/monkey_island/cc/services/reset_agent_configuration.py

@@ -0,0 +1,25 @@
+from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
+
+
+class reset_agent_configuration:
+    def __init__(
+        self,
+        agent_configuration_repository: IAgentConfigurationRepository,
+        file_repository: IFileRepository,
+    ):
+        self._agent_configuration_repository = agent_configuration_repository
+        self._file_repository = file_repository
+
+    def __call__(self):
+        self._remove_pba_files()
+        self._agent_configuration_repository.reset_to_default()
+
+    def _remove_pba_files(self):
+        agent_configuration = self._agent_configuration_repository.get_configuration()
+        custom_pbas = agent_configuration.custom_pbas
+
+        if custom_pbas.linux_filename:
+            self._file_repository.delete_file(custom_pbas.linux_filename)
+
+        if custom_pbas.windows_filename:
+            self._file_repository.delete_file(custom_pbas.windows_filename)

monkey/tests/unit_tests/monkey_island/cc/services/test_reset_agent_configuration.py

@@ -0,0 +1,56 @@
+from unittest.mock import MagicMock
+
+import pytest
+from tests.monkey_island import InMemoryAgentConfigurationRepository
+
+from common.agent_configuration import AgentConfiguration
+from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
+from monkey_island.cc.services import reset_agent_configuration
+
+LINUX_FILENAME = "linux_pba_file.sh"
+WINDOWS_FILENAME = "windows_pba_file.ps1"
+
+
+@pytest.fixture
+def agent_configuration(default_agent_configuration: AgentConfiguration) -> AgentConfiguration:
+    custom_pbas = default_agent_configuration.custom_pbas.copy(
+        update={"linux_filename": LINUX_FILENAME, "windows_filename": WINDOWS_FILENAME},
+    )
+    return default_agent_configuration.copy(update={"custom_pbas": custom_pbas})
+
+
+@pytest.fixture
+def agent_configuration_repository(
+    agent_configuration: AgentConfiguration,
+) -> IAgentConfigurationRepository:
+    agent_configuration_repository = InMemoryAgentConfigurationRepository()
+    agent_configuration_repository.store_configuration(agent_configuration)
+
+    return agent_configuration_repository
+
+
+@pytest.fixture
+def mock_file_repository() -> IFileRepository:
+    return MagicMock(spec=IFileRepository)
+
+
+@pytest.fixture
+def callable_reset_agent_configuration(agent_configuration_repository, mock_file_repository):
+    return reset_agent_configuration(agent_configuration_repository, mock_file_repository)
+
+
+def test_reset_configuration__remove_pba_files(
+    callable_reset_agent_configuration, mock_file_repository
+):
+    callable_reset_agent_configuration()
+
+    assert mock_file_repository.delete_file.called_with(LINUX_FILENAME)
+    assert mock_file_repository.delete_file.called_with(WINDOWS_FILENAME)
+
+
+def test_reset_configuration__agent_configuration_changed(
+    callable_reset_agent_configuration, agent_configuration_repository, agent_configuration
+):
+    callable_reset_agent_configuration()
+
+    assert agent_configuration_repository.get_configuration() != agent_configuration