From fc76bb2d83869070766e68bd306d81c25d9d3d38 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:04:48 -0400 Subject: [PATCH] Island: Set config on mode set in IslandModeService --- monkey/common/configuration/__init__.py | 1 + .../default_agent_configuration.py | 6 ++++ .../monkey_island/cc/services/initialize.py | 11 ++++++- .../cc/services/mode/island_mode_service.py | 25 +++++++++++++-- .../cc/resources/test_island_mode.py | 17 +++++++--- .../cc/services/test_island_mode_service.py | 32 +++++++++++++++++-- 6 files changed, 82 insertions(+), 10 deletions(-) diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index fc1f3c84d..fb52b54e2 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -12,4 +12,5 @@ from .agent_sub_configurations import ( ) from .default_agent_configuration import ( DEFAULT_AGENT_CONFIGURATION, + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, ) diff --git a/monkey/common/configuration/default_agent_configuration.py b/monkey/common/configuration/default_agent_configuration.py index 251676017..9045429ca 100644 --- a/monkey/common/configuration/default_agent_configuration.py +++ b/monkey/common/configuration/default_agent_configuration.py @@ -1,3 +1,5 @@ +import dataclasses + from . import AgentConfiguration from .agent_sub_configurations import ( CustomPBAConfiguration, @@ -112,3 +114,7 @@ DEFAULT_AGENT_CONFIGURATION = AgentConfiguration( payloads=PAYLOAD_CONFIGURATION, propagation=PROPAGATION_CONFIGURATION, ) + +DEFAULT_RANSOMWARE_AGENT_CONFIGURATION = dataclasses.replace( + DEFAULT_AGENT_CONFIGURATION, post_breach_actions=[] +) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index fafd274d3..30301b90f 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -3,7 +3,11 @@ from pathlib import Path from common import DIContainer from common.aws import AWSInstance -from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfiguration +from common.configuration import ( + DEFAULT_AGENT_CONFIGURATION, + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, + AgentConfiguration, +) from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, @@ -36,6 +40,11 @@ def initialize_services(data_dir: Path) -> DIContainer: container.register_convention( AgentConfiguration, "default_agent_configuration", DEFAULT_AGENT_CONFIGURATION ) + container.register_convention( + AgentConfiguration, + "default_ransomware_agent_configuration", + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, + ) container.register_instance(AWSInstance, AWSInstance()) container.register_instance( diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py index dd67b6643..5183bdb8c 100644 --- a/monkey/monkey_island/cc/services/mode/island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -1,10 +1,20 @@ -from monkey_island.cc.repository import ISimulationRepository +from common.configuration import AgentConfiguration +from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class IslandModeService: - def __init__(self, simulation_repository: ISimulationRepository): + def __init__( + self, + _agent_configuration_repository: IAgentConfigurationRepository, + simulation_repository: ISimulationRepository, + default_agent_configuration: AgentConfiguration, + default_ransomware_agent_configuration: AgentConfiguration, + ): + self._agent_configuration_repository = _agent_configuration_repository self._simulation_repository = simulation_repository + self._default_agent_configuration = default_agent_configuration + self._default_ransomware_agent_configuration = default_ransomware_agent_configuration def get_mode(self): """ @@ -23,3 +33,14 @@ class IslandModeService: :raises StorageError: If the mode could not be saved """ self._simulation_repository.set_mode(mode) + self._set_configuration(mode) + + def _set_configuration(self, mode: IslandModeEnum): + if mode == IslandModeEnum.RANSOMWARE: + self._agent_configuration_repository.store_configuration( + self._default_ransomware_agent_configuration + ) + else: + self._agent_configuration_repository.store_configuration( + self._default_agent_configuration + ) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index fec2b6d37..ae3905923 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py @@ -5,18 +5,27 @@ import pytest from tests.common import StubDIContainer from tests.monkey_island import InMemorySimulationRepository -from monkey_island.cc.repository import ISimulationRepository, RetrievalError +from monkey_island.cc.repository import RetrievalError from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource from monkey_island.cc.services import IslandModeService from monkey_island.cc.services.mode.mode_enum import IslandModeEnum +class MockIslandModeService(IslandModeService): + def __init__(self): + self._simulation_repository = InMemorySimulationRepository() + + def get_mode(self) -> IslandModeEnum: + return self._simulation_repository.get_mode() + + def set_mode(self, mode: IslandModeEnum): + self._simulation_repository.set_mode(mode) + + @pytest.fixture def flask_client(build_flask_client): container = StubDIContainer() - - container.register(ISimulationRepository, InMemorySimulationRepository) - container.register_instance(IslandModeService, container.resolve(IslandModeService)) + container.register_instance(IslandModeService, MockIslandModeService()) with build_flask_client(container) as flask_client: yield flask_client diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py index 79074b213..4d5f5f4bc 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py @@ -1,16 +1,42 @@ import pytest -from tests.monkey_island import InMemorySimulationRepository +from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository +from common.configuration import DEFAULT_AGENT_CONFIGURATION, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION from monkey_island.cc.services import IslandModeService from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @pytest.fixture -def island_mode_service(): - return IslandModeService(InMemorySimulationRepository()) +def agent_configuration_repository(): + return InMemoryAgentConfigurationRepository() + + +@pytest.fixture +def island_mode_service(agent_configuration_repository): + return IslandModeService( + agent_configuration_repository, + InMemorySimulationRepository(), + DEFAULT_AGENT_CONFIGURATION, + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, + ) @pytest.mark.parametrize("mode", list(IslandModeEnum)) def test_set_mode(island_mode_service, mode): island_mode_service.set_mode(mode) assert island_mode_service.get_mode() == mode + + +@pytest.mark.parametrize( + "mode, expected_config", + [ + (IslandModeEnum.UNSET, DEFAULT_AGENT_CONFIGURATION), + (IslandModeEnum.ADVANCED, DEFAULT_AGENT_CONFIGURATION), + (IslandModeEnum.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION), + ], +) +def test_set_mode_sets_config( + island_mode_service, agent_configuration_repository, mode, expected_config +): + island_mode_service.set_mode(mode) + assert agent_configuration_repository.get_configuration() == expected_config