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,
AgentConfigurationSchema,
)
from monkey_island.cc import repository
from monkey_island.cc.repository import (
IAgentConfigurationRepository,
IFileRepository,
@ -25,10 +26,10 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository):
configuration_json = f.read().decode()
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:
except repository.FileNotFoundError:
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):
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.monkey_island import SingleFileRepository
from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository
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():
@ -24,3 +25,10 @@ def test_get_default_agent_config():
retrieved_agent_configuration = repository.get_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()