diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 885d467be..ffe00e736 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -1,3 +1,5 @@ from .file_storage import FileRetrievalError, IFileRepository, LocalStorageFileRepository from .i_agent_binary_repository import IAgentBinaryRepository, AgentRetrievalError from .agent_binary_repository import AgentBinaryRepository +from .i_agent_configuration_repository import IAgentConfigurationRepository +from .file_agent_configuration_repository import FileAgentConfigurationRepository diff --git a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py new file mode 100644 index 000000000..e9724aa4b --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py @@ -0,0 +1,25 @@ +import io + +from common.configuration import AgentConfiguration, AgentConfigurationSchema +from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository + +AGENT_CONFIGURATION_FILE_NAME = "agent_configuration.json" + + +class FileAgentConfigurationRepository(IAgentConfigurationRepository): + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + 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() + + return self._schema.loads(configuration_json) + + def set_configuration(self, agent_configuration: AgentConfiguration): + configuration_json = self._schema.dumps(agent_configuration) + + self._file_repository.save_file( + AGENT_CONFIGURATION_FILE_NAME, io.BytesIO(configuration_json.encode()) + ) 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 new file mode 100644 index 000000000..bfb894913 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py @@ -0,0 +1,16 @@ +from tests.common.example_agent_configuration import AGENT_CONFIGURATION +from tests.monkey_island import SingleFileRepository + +from common.configuration import AgentConfigurationSchema +from monkey_island.cc.repository import FileAgentConfigurationRepository + + +def test_store_agent_config(): + repository = FileAgentConfigurationRepository(SingleFileRepository()) + schema = AgentConfigurationSchema() + agent_configuration = schema.load(AGENT_CONFIGURATION) + + repository.set_configuration(agent_configuration) + retrieved_agent_configuration = repository.get_configuration() + + assert retrieved_agent_configuration == agent_configuration