Island: Return default configuration from get_configuration()

This commit is contained in:
Mike Salvatore 2022-06-17 12:55:28 -04:00
parent a3e3e3e324
commit 45168b5ba7
3 changed files with 31 additions and 7 deletions

View File

@ -1,7 +1,15 @@
import io import io
from common.configuration import AgentConfiguration, AgentConfigurationSchema from common.configuration import (
from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository DEFAULT_AGENT_CONFIGURATION,
AgentConfiguration,
AgentConfigurationSchema,
)
from monkey_island.cc.repository import (
IAgentConfigurationRepository,
IFileRepository,
RetrievalError,
)
AGENT_CONFIGURATION_FILE_NAME = "agent_configuration.json" AGENT_CONFIGURATION_FILE_NAME = "agent_configuration.json"
@ -12,10 +20,15 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository):
self._schema = AgentConfigurationSchema() self._schema = AgentConfigurationSchema()
def get_configuration(self) -> AgentConfiguration: def get_configuration(self) -> AgentConfiguration:
with self._file_repository.open_file(AGENT_CONFIGURATION_FILE_NAME) as f: try:
configuration_json = f.read().decode() 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): def store_configuration(self, agent_configuration: AgentConfiguration):
configuration_json = self._schema.dumps(agent_configuration) configuration_json = self._schema.dumps(agent_configuration)

View File

@ -13,7 +13,8 @@ class IAgentConfigurationRepository(ABC):
""" """
Retrieve the agent configuration from the repository 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 :raises RetrievalError: if the configuration can not be retrieved
""" """
pass pass

View File

@ -1,7 +1,7 @@
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 SingleFileRepository
from common.configuration import AgentConfigurationSchema from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema
from monkey_island.cc.repository import FileAgentConfigurationRepository from monkey_island.cc.repository import FileAgentConfigurationRepository
@ -14,3 +14,13 @@ def test_store_agent_config():
retrieved_agent_configuration = repository.get_configuration() retrieved_agent_configuration = repository.get_configuration()
assert retrieved_agent_configuration == agent_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