From 1f00a13649d784d80dc89d9685799c23bc41df82 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 22 Jun 2022 07:35:06 -0400 Subject: [PATCH] Island: Pass default agent configuration to repository constructor --- monkey/common/configuration/__init__.py | 5 ++++- .../default_agent_configuration.py | 9 +++++++- .../file_agent_configuration_repository.py | 13 +++++------ .../monkey_island/cc/services/initialize.py | 4 ++++ .../common/test_agent_configuration.py | 4 ++-- monkey/tests/unit_tests/conftest.py | 7 ++++++ ...est_file_agent_configuration_repository.py | 22 ++++++++++--------- 7 files changed, 43 insertions(+), 21 deletions(-) diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index 7a9131cc7..2246e27f2 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -2,4 +2,7 @@ from .agent_configuration import ( AgentConfiguration, AgentConfigurationSchema, ) -from .default_agent_configuration import DEFAULT_AGENT_CONFIGURATION +from .default_agent_configuration import ( + DEFAULT_AGENT_CONFIGURATION_JSON, + build_default_agent_configuration, +) diff --git a/monkey/common/configuration/default_agent_configuration.py b/monkey/common/configuration/default_agent_configuration.py index 933b723ea..8c909f15d 100644 --- a/monkey/common/configuration/default_agent_configuration.py +++ b/monkey/common/configuration/default_agent_configuration.py @@ -1,4 +1,6 @@ -DEFAULT_AGENT_CONFIGURATION = """{ +from . import AgentConfiguration, AgentConfigurationSchema + +DEFAULT_AGENT_CONFIGURATION_JSON = """{ "keep_tunnel_open_time": 30, "post_breach_actions": [ { @@ -215,3 +217,8 @@ DEFAULT_AGENT_CONFIGURATION = """{ } } """ + + +def build_default_agent_configuration() -> AgentConfiguration: + schema = AgentConfigurationSchema() + return schema.loads(DEFAULT_AGENT_CONFIGURATION_JSON) diff --git a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py index 89fc57960..b63ee817c 100644 --- a/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py +++ b/monkey/monkey_island/cc/repository/file_agent_configuration_repository.py @@ -1,10 +1,6 @@ import io -from common.configuration import ( - DEFAULT_AGENT_CONFIGURATION, - AgentConfiguration, - AgentConfigurationSchema, -) +from common.configuration import AgentConfiguration, AgentConfigurationSchema from monkey_island.cc import repository from monkey_island.cc.repository import ( IAgentConfigurationRepository, @@ -16,7 +12,10 @@ AGENT_CONFIGURATION_FILE_NAME = "agent_configuration.json" class FileAgentConfigurationRepository(IAgentConfigurationRepository): - def __init__(self, file_repository: IFileRepository): + def __init__( + self, default_agent_configuration: AgentConfiguration, file_repository: IFileRepository + ): + self._default_agent_configuration = default_agent_configuration self._file_repository = file_repository self._schema = AgentConfigurationSchema() @@ -27,7 +26,7 @@ class FileAgentConfigurationRepository(IAgentConfigurationRepository): return self._schema.loads(configuration_json) except repository.FileNotFoundError: - return self._schema.loads(DEFAULT_AGENT_CONFIGURATION) + return self._default_agent_configuration except Exception as err: raise RetrievalError(f"Error retrieving the agent configuration: {err}") diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 4b03daf7d..52343bbf5 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -3,6 +3,7 @@ from pathlib import Path from common import DIContainer from common.aws import AWSInstance +from common.configuration import AgentConfiguration, build_default_agent_configuration from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, @@ -30,6 +31,9 @@ def initialize_services(data_dir: Path) -> DIContainer: container = DIContainer() container.register_convention(Path, "data_dir", data_dir) + container.register_convention( + AgentConfiguration, "default_agent_configuration", build_default_agent_configuration() + ) container.register_instance(AWSInstance, AWSInstance()) container.register_instance( diff --git a/monkey/tests/unit_tests/common/test_agent_configuration.py b/monkey/tests/unit_tests/common/test_agent_configuration.py index 733e2e708..3383ab1b7 100644 --- a/monkey/tests/unit_tests/common/test_agent_configuration.py +++ b/monkey/tests/unit_tests/common/test_agent_configuration.py @@ -25,7 +25,7 @@ from tests.common.example_agent_configuration import ( from common import OperatingSystems from common.configuration import ( - DEFAULT_AGENT_CONFIGURATION, + DEFAULT_AGENT_CONFIGURATION_JSON, AgentConfiguration, AgentConfigurationSchema, ) @@ -180,6 +180,6 @@ def test_agent_configuration(): def test_default_agent_configuration(): schema = AgentConfigurationSchema() - config = schema.loads(DEFAULT_AGENT_CONFIGURATION) + config = schema.loads(DEFAULT_AGENT_CONFIGURATION_JSON) assert isinstance(config, AgentConfiguration) diff --git a/monkey/tests/unit_tests/conftest.py b/monkey/tests/unit_tests/conftest.py index 5b83a7e69..2f00d3d16 100644 --- a/monkey/tests/unit_tests/conftest.py +++ b/monkey/tests/unit_tests/conftest.py @@ -6,6 +6,8 @@ from typing import Callable, Dict import pytest from _pytest.monkeypatch import MonkeyPatch +from common.configuration import AgentConfiguration, build_default_agent_configuration + MONKEY_BASE_PATH = str(Path(__file__).parent.parent.parent) sys.path.insert(0, MONKEY_BASE_PATH) @@ -54,3 +56,8 @@ def load_monkey_config(data_for_tests_dir) -> Callable[[str], Dict]: return json.loads(open(config_path, "r").read()) return inner + + +@pytest.fixture +def default_agent_configuration() -> AgentConfiguration: + return build_default_agent_configuration() diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py index 93a80dea6..4ab111606 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_agent_configuration_repository.py @@ -2,12 +2,16 @@ import pytest from tests.common.example_agent_configuration import AGENT_CONFIGURATION from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfigurationSchema +from common.configuration import AgentConfigurationSchema from monkey_island.cc.repository import FileAgentConfigurationRepository, RetrievalError -def test_store_agent_config(): - repository = FileAgentConfigurationRepository(SingleFileRepository()) +@pytest.fixture +def repository(default_agent_configuration): + return FileAgentConfigurationRepository(default_agent_configuration, SingleFileRepository()) + + +def test_store_agent_config(repository): schema = AgentConfigurationSchema() agent_configuration = schema.load(AGENT_CONFIGURATION) @@ -17,18 +21,16 @@ def test_store_agent_config(): 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) - +def test_get_default_agent_config(repository, default_agent_configuration): retrieved_agent_configuration = repository.get_configuration() assert retrieved_agent_configuration == default_agent_configuration -def test_get_agent_config_retrieval_error(): - repository = FileAgentConfigurationRepository(OpenErrorFileRepository()) +def test_get_agent_config_retrieval_error(default_agent_configuration): + repository = FileAgentConfigurationRepository( + default_agent_configuration, OpenErrorFileRepository() + ) with pytest.raises(RetrievalError): repository.get_configuration()