forked from p15670423/monkey
Island: Return default configuration from get_configuration()
This commit is contained in:
parent
a3e3e3e324
commit
45168b5ba7
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue