Island: Set config on mode set in IslandModeService

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,
 )

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=[]
+)

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(

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
+            )

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

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