Island: Raise RetrievalError in FileAgentConfigurationRepository

This commit is contained in:
Mike Salvatore 2022-06-21 09:40:00 -04:00
parent 4de9f3cb6d
commit c008db4cf2
2 changed files with 14 additions and 5 deletions

View File

@ -5,6 +5,7 @@ from common.configuration import (
AgentConfiguration, AgentConfiguration,
AgentConfigurationSchema, AgentConfigurationSchema,
) )
from monkey_island.cc import repository
from monkey_island.cc.repository import ( from monkey_island.cc.repository import (
IAgentConfigurationRepository, IAgentConfigurationRepository,
IFileRepository, IFileRepository,
@ -25,10 +26,10 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository):
configuration_json = f.read().decode() configuration_json = f.read().decode()
return self._schema.loads(configuration_json) return self._schema.loads(configuration_json)
# TODO: Handle FileRetrievalError vs FileNotFoundError except repository.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) return self._schema.loads(DEFAULT_AGENT_CONFIGURATION)
except Exception as err:
raise RetrievalError(f"Error retrieving the agent configuration: {err}")
def store_configuration(self, agent_configuration: AgentConfiguration): def store_configuration(self, agent_configuration: AgentConfiguration):
configuration_json = self._schema.dumps(agent_configuration) configuration_json = self._schema.dumps(agent_configuration)

View File

@ -1,8 +1,9 @@
import pytest
from tests.common.example_agent_configuration import AGENT_CONFIGURATION from tests.common.example_agent_configuration import AGENT_CONFIGURATION
from tests.monkey_island import SingleFileRepository from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository
from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema
from monkey_island.cc.repository import FileAgentConfigurationRepository from monkey_island.cc.repository import FileAgentConfigurationRepository, RetrievalError
def test_store_agent_config(): def test_store_agent_config():
@ -24,3 +25,10 @@ def test_get_default_agent_config():
retrieved_agent_configuration = repository.get_configuration() retrieved_agent_configuration = repository.get_configuration()
assert retrieved_agent_configuration == default_agent_configuration assert retrieved_agent_configuration == default_agent_configuration
def test_get_agent_config_retrieval_error():
repository = FileAgentConfigurationRepository(OpenErrorFileRepository())
with pytest.raises(RetrievalError):
repository.get_configuration()