From aa7509e49f95811cc96ed3eb230a2ac3495a1906 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 30 Jun 2022 10:02:01 -0400 Subject: [PATCH 1/3] Island: Add reset_to_default() to IAgentConfigurationRepository --- .../cc/repository/i_agent_configuration_repository.py | 9 +++++++++ 1 file changed, 9 insertions(+) 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 From e293639a22e34bc67d6be564d62db9d8d0b6532e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 30 Jun 2022 10:03:21 -0400 Subject: [PATCH 2/3] Island: Add reset_to_default() to FileAgentConfigurationRepository --- .../repository/file_agent_configuration_repository.py | 3 +++ .../test_file_agent_configuration_repository.py | 10 ++++++++++ 2 files changed, 13 insertions(+) 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/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 From 41b1c56fd06fa2a5f270803daddbb7909b3df5ad Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 30 Jun 2022 11:35:29 -0400 Subject: [PATCH 3/3] UT: Add reset_to_default() to InMemoryAgentConfigurationRepository --- .../in_memory_agent_configuration_repository.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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