forked from p15670423/monkey
Island: Raise RetrievalError in FileAgentConfigurationRepository
This commit is contained in:
parent
4de9f3cb6d
commit
c008db4cf2
|
@ -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)
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue