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 2cd2c8c93..fcccd49a1 100644 --- a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py @@ -1,7 +1,15 @@ import io -from common.configuration import AgentConfiguration, AgentConfigurationSchema -from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository +from common.configuration import ( + DEFAULT_AGENT_CONFIGURATION, + AgentConfiguration, + AgentConfigurationSchema, +) +from monkey_island.cc.repository import ( + IAgentConfigurationRepository, + IFileRepository, + RetrievalError, +) AGENT_CONFIGURATION_FILE_NAME = "agent_configuration.json" @@ -12,10 +20,15 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository): self._schema = AgentConfigurationSchema() def get_configuration(self) -> AgentConfiguration: - with self._file_repository.open_file(AGENT_CONFIGURATION_FILE_NAME) as f: - configuration_json = f.read().decode() + try: + with self._file_repository.open_file(AGENT_CONFIGURATION_FILE_NAME) as f: + configuration_json = f.read().decode() - return self._schema.loads(configuration_json) + return self._schema.loads(configuration_json) + # TODO: Handle FileRetrievalError vs FileNotFoundError + # https://github.com/guardicore/monkey/blob/e8001d8cf76340e42bf17ff62523bd2d85fc4841/monkey/monkey_island/cc/repository/file_storage/local_storage_file_repository.py#L47-L50 + except RetrievalError: + return self._schema.loads(DEFAULT_AGENT_CONFIGURATION) def store_configuration(self, agent_configuration: AgentConfiguration): configuration_json = self._schema.dumps(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 d99934ad9..105d26fe9 100644 --- a/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/i_agent_configuration_repository.py @@ -13,7 +13,8 @@ class IAgentConfigurationRepository(ABC): """ Retrieve the agent configuration from the repository - :return: The agent configuration + :return: The agent configuration as retrieved from the repository, or the default + configuration if none could be retrieved. :raises RetrievalError: if the configuration can not be retrieved """ pass 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 941dd2586..7ad066623 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 @@ -1,7 +1,7 @@ from tests.common.example_agent_configuration import AGENT_CONFIGURATION from tests.monkey_island import SingleFileRepository -from common.configuration import AgentConfigurationSchema +from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema from monkey_island.cc.repository import FileAgentConfigurationRepository @@ -14,3 +14,13 @@ def test_store_agent_config(): retrieved_agent_configuration = repository.get_configuration() assert retrieved_agent_configuration == agent_configuration + + +def test_get_default_agent_config(): + repository = FileAgentConfigurationRepository(SingleFileRepository()) + schema = AgentConfigurationSchema() + default_agent_configuration = schema.loads(DEFAULT_AGENT_CONFIGURATION) + + retrieved_agent_configuration = repository.get_configuration() + + assert retrieved_agent_configuration == default_agent_configuration