diff --git a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py index 312e3921e..980697657 100644 --- a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py @@ -35,3 +35,6 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository): self._file_repository.save_file( AGENT_CONFIGURATION_FILE_NAME, io.BytesIO(configuration_json.encode()) ) + + def reset_to_default(self): + self.store_configuration(self._default_agent_configuration) diff --git a/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py index ef63121c2..d941ce075 100644 --- a/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py @@ -28,3 +28,12 @@ class IAgentConfigurationRepository(ABC): :raises StorageError: If the configuration could not be stored """ pass + + @abstractmethod + def reset_to_default(self): + """ + Remove any stored configuration from the repository + + :raises RemovalError: If the repository could not be reset + """ + pass diff --git a/monkey/tests/monkey_island/in_memory_agent_configuration_repository.py b/monkey/tests/monkey_island/in_memory_agent_configuration_repository.py index e9bcbae62..cf97ac08a 100644 --- a/monkey/tests/monkey_island/in_memory_agent_configuration_repository.py +++ b/monkey/tests/monkey_island/in_memory_agent_configuration_repository.py @@ -6,10 +6,14 @@ from monkey_island.cc.repository import IAgentConfigurationRepository class InMemoryAgentConfigurationRepository(IAgentConfigurationRepository): def __init__(self): - self._configuration = AgentConfiguration.from_mapping(AGENT_CONFIGURATION) + self._default_configuration = AgentConfiguration.from_mapping(AGENT_CONFIGURATION) + self._configuration = self._default_configuration def get_configuration(self): return self._configuration def store_configuration(self, agent_configuration): self._configuration = agent_configuration + + def reset_to_default(self): + self._configuration = self._default_configuration diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py index fb7863dc3..1afde1f6f 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py @@ -33,3 +33,13 @@ def test_get_agent_config_retrieval_error(default_agent_configuration): with pytest.raises(RetrievalError): repository.get_configuration() + + +def test_reset_to_default(repository, default_agent_configuration): + agent_configuration = AgentConfiguration.from_mapping(AGENT_CONFIGURATION) + + repository.store_configuration(agent_configuration) + repository.reset_to_default() + retrieved_agent_configuration = repository.get_configuration() + + assert retrieved_agent_configuration == default_agent_configuration