Merge pull request #2327 from guardicore/2266-replace-IslandModeService

Replace IslandModeService with event

monkey/monkey_island/cc/event_queue/i_island_event_queue.py

@@ -9,6 +9,7 @@ class IslandEventTopic(Enum):
     AGENT_CONNECTED = auto()
     CLEAR_SIMULATION_DATA = auto()
     RESET_AGENT_CONFIGURATION = auto()
+    SET_ISLAND_MODE = auto()
 
 
 class IIslandEventQueue(ABC):

monkey/monkey_island/cc/island_event_handlers/__init__.py

@@ -1,2 +1,3 @@
 from .reset_agent_configuration import reset_agent_configuration
 from .reset_machine_repository import reset_machine_repository
+from .set_agent_configuration_per_island_mode import set_agent_configuration_per_island_mode

monkey/monkey_island/cc/island_event_handlers/set_agent_configuration_per_island_mode.py

@@ -0,0 +1,29 @@
+from common.agent_configuration import AgentConfiguration
+from monkey_island.cc.models import IslandMode
+from monkey_island.cc.repository import IAgentConfigurationRepository
+
+
+class set_agent_configuration_per_island_mode:
+    """
+    Callable class that sets the default Agent configuration as per the Island's mode
+    """
+
+    def __init__(
+        self,
+        agent_configuration_repository: IAgentConfigurationRepository,
+        default_agent_configuration: AgentConfiguration,
+        default_ransomware_agent_configuration: AgentConfiguration,
+    ):
+        self._agent_configuration_repository = agent_configuration_repository
+        self._default_agent_configuration = default_agent_configuration
+        self._default_ransomware_agent_configuration = default_ransomware_agent_configuration
+
+    def __call__(self, mode: IslandMode):
+        if mode == IslandMode.RANSOMWARE:
+            self._agent_configuration_repository.store_configuration(
+                self._default_ransomware_agent_configuration
+            )
+        else:
+            self._agent_configuration_repository.store_configuration(
+                self._default_agent_configuration
+            )

monkey/monkey_island/cc/resources/island_mode.py

@@ -4,10 +4,11 @@ from http import HTTPStatus
 
 from flask import request
 
+from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.models import IslandMode as IslandModeEnum
+from monkey_island.cc.repository import ISimulationRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
-from monkey_island.cc.services import IslandModeService
 
 logger = logging.getLogger(__name__)
 
@@ -15,16 +16,19 @@ logger = logging.getLogger(__name__)
 class IslandMode(AbstractResource):
     urls = ["/api/island/mode"]
 
-    def __init__(self, island_mode_service: IslandModeService):
+    def __init__(
-        self._island_mode_service = island_mode_service
+        self,
+        island_event_queue: IIslandEventQueue,
+        simulation_repository: ISimulationRepository,
+    ):
+        self._island_event_queue = island_event_queue
+        self._simulation_repository = simulation_repository
 
     @jwt_required
     def put(self):
         try:
             mode = IslandModeEnum(request.json)
-
+            self._island_event_queue.publish(topic=IslandEventTopic.SET_ISLAND_MODE, mode=mode)
-            self._island_mode_service.set_mode(mode)
-
             return {}, HTTPStatus.NO_CONTENT
         except (AttributeError, json.decoder.JSONDecodeError):
             return {}, HTTPStatus.BAD_REQUEST
@@ -33,5 +37,5 @@ class IslandMode(AbstractResource):
 
     @jwt_required
     def get(self):
-        island_mode = self._island_mode_service.get_mode()
+        island_mode = self._simulation_repository.get_mode()
         return island_mode.value, HTTPStatus.OK

monkey/monkey_island/cc/services/__init__.py

@@ -1,4 +1,3 @@
 from .authentication_service import AuthenticationService
 
 from .aws import AWSService
-from .island_mode_service import IslandModeService

monkey/monkey_island/cc/services/initialize.py

@@ -47,7 +47,7 @@ from monkey_island.cc.repository import (
 )
 from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
 from monkey_island.cc.server_utils.encryption import ILockableEncryptor, RepositoryEncryptor
-from monkey_island.cc.services import AWSService, IslandModeService
+from monkey_island.cc.services import AWSService
 from monkey_island.cc.services.attack.technique_reports.T1003 import T1003, T1003GetReportData
 from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
 from monkey_island.cc.setup.mongo.mongo_setup import MONGO_URL
@@ -175,7 +175,6 @@ def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository):
 def _register_services(container: DIContainer):
     container.register_instance(AWSService, container.resolve(AWSService))
     container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService))
-    container.register_instance(IslandModeService, container.resolve(IslandModeService))
     container.register_instance(AuthenticationService, container.resolve(AuthenticationService))
 
 

monkey/monkey_island/cc/services/island_mode_service.py

@@ -1,46 +0,0 @@
-from common.agent_configuration import AgentConfiguration
-from monkey_island.cc.models import IslandMode
-from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository
-
-
-class IslandModeService:
-    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):
-        """
-        Get's the island's current mode
-
-        :return The island's current mode
-        :raises RetrievalError: If the mode could not be retrieved
-        """
-        return self._simulation_repository.get_mode()
-
-    def set_mode(self, mode: IslandMode):
-        """
-        Set the island's mode
-
-        :param mode: The island's new mode
-        :raises StorageError: If the mode could not be saved
-        """
-        self._simulation_repository.set_mode(mode)
-        self._set_configuration(mode)
-
-    def _set_configuration(self, mode: IslandMode):
-        if mode == IslandMode.RANSOMWARE:
-            self._agent_configuration_repository.store_configuration(
-                self._default_ransomware_agent_configuration
-            )
-        else:
-            self._agent_configuration_repository.store_configuration(
-                self._default_agent_configuration
-            )

monkey/monkey_island/cc/setup/island_event_handlers.py

@@ -5,12 +5,14 @@ from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.island_event_handlers import (
     reset_agent_configuration,
     reset_machine_repository,
+    set_agent_configuration_per_island_mode,
 )
 from monkey_island.cc.repository import (
     IAgentEventRepository,
     IAgentRepository,
     ICredentialsRepository,
     INodeRepository,
+    ISimulationRepository,
 )
 from monkey_island.cc.services.database import Database
 
@@ -20,30 +22,29 @@ def setup_island_event_handlers(container: DIContainer):
 
     _subscribe_reset_agent_configuration_events(island_event_queue, container)
     _subscribe_clear_simulation_data_events(island_event_queue, container)
+    _subscribe_set_island_mode_events(island_event_queue, container)
 
 
 def _subscribe_reset_agent_configuration_events(
     island_event_queue: IIslandEventQueue, container: DIContainer
 ):
-    island_event_queue.subscribe(
+    topic = IslandEventTopic.RESET_AGENT_CONFIGURATION
-        IslandEventTopic.RESET_AGENT_CONFIGURATION, container.resolve(reset_agent_configuration)
+
-    )
+    island_event_queue.subscribe(topic, container.resolve(reset_agent_configuration))
 
 
 def _subscribe_clear_simulation_data_events(
     island_event_queue: IIslandEventQueue, container: DIContainer
 ):
+    topic = IslandEventTopic.CLEAR_SIMULATION_DATA
+
     legacy_database_reset = partial(Database.reset_db, reset_config=False)
-    island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, legacy_database_reset)
+    island_event_queue.subscribe(topic, legacy_database_reset)
 
     credentials_repository = container.resolve(ICredentialsRepository)
-    island_event_queue.subscribe(
+    island_event_queue.subscribe(topic, credentials_repository.remove_stolen_credentials)
-        IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials
-    )
 
-    island_event_queue.subscribe(
+    island_event_queue.subscribe(topic, container.resolve(reset_machine_repository))
-        IslandEventTopic.CLEAR_SIMULATION_DATA, container.resolve(reset_machine_repository)
-    )
 
     for i_repository in [
         INodeRepository,
@@ -51,4 +52,15 @@ def _subscribe_clear_simulation_data_events(
         IAgentRepository,
     ]:
         repository = container.resolve(i_repository)
-        island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, repository.reset)
+        island_event_queue.subscribe(topic, repository.reset)
+
+
+def _subscribe_set_island_mode_events(
+    island_event_queue: IIslandEventQueue, container: DIContainer
+):
+    topic = IslandEventTopic.SET_ISLAND_MODE
+
+    island_event_queue.subscribe(topic, container.resolve(set_agent_configuration_per_island_mode))
+
+    simulation_repository = container.resolve(ISimulationRepository)
+    island_event_queue.subscribe(topic, simulation_repository.set_mode)

monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py

@@ -5,30 +5,42 @@ import pytest
 from tests.common import StubDIContainer
 from tests.monkey_island import InMemorySimulationRepository
 
+from monkey_island.cc.event_queue import IIslandEventQueue
 from monkey_island.cc.models import IslandMode
-from monkey_island.cc.repository import RetrievalError
+from monkey_island.cc.repository import ISimulationRepository
 from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource
-from monkey_island.cc.services import IslandModeService
-
-
-class MockIslandModeService(IslandModeService):
-    def __init__(self):
-        self._simulation_repository = InMemorySimulationRepository()
-
-    def get_mode(self) -> IslandMode:
-        return self._simulation_repository.get_mode()
-
-    def set_mode(self, mode: IslandMode):
-        self._simulation_repository.set_mode(mode)
 
 
 @pytest.fixture
-def flask_client(build_flask_client):
+def flask_client_builder(build_flask_client):
+    def inner(side_effect=None):
         container = StubDIContainer()
-    container.register_instance(IslandModeService, MockIslandModeService())
+
+        in_memory_simulation_repository = InMemorySimulationRepository()
+        container.register_instance(ISimulationRepository, in_memory_simulation_repository)
+
+        mock_island_event_queue = MagicMock(spec=IIslandEventQueue)
+        mock_island_event_queue.publish.side_effect = (
+            side_effect
+            if side_effect
+            else lambda topic, mode: in_memory_simulation_repository.set_mode(mode)
+        )
+        container.register_instance(IIslandEventQueue, mock_island_event_queue)
 
         with build_flask_client(container) as flask_client:
-        yield flask_client
+            return flask_client
+
+    return inner
+
+
+@pytest.fixture
+def flask_client(flask_client_builder):
+    return flask_client_builder()
+
+
+@pytest.fixture
+def flask_client__internal_server_error(flask_client_builder):
+    return flask_client_builder(Exception)
 
 
 @pytest.mark.parametrize(
@@ -53,20 +65,12 @@ def test_island_mode_post__invalid_mode(flask_client):
     assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
 
 
-def test_island_mode_post__internal_server_error(build_flask_client):
+def test_island_mode_post__internal_server_error(flask_client__internal_server_error):
-    mock_island_mode_service = MagicMock(spec=IslandModeService)
+    resp = flask_client__internal_server_error.put(
-    mock_island_mode_service.set_mode = MagicMock(side_effect=RetrievalError)
-
-    container = StubDIContainer()
-    container.register_instance(IslandModeService, mock_island_mode_service)
-
-    with build_flask_client(container) as flask_client:
-        resp = flask_client.put(
         IslandModeResource.urls[0],
         json=IslandMode.RANSOMWARE.value,
         follow_redirects=True,
     )
-
     assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
 
 

monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py

@@ -1,45 +0,0 @@
-import pytest
-from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository
-
-from common.agent_configuration import (
-    DEFAULT_AGENT_CONFIGURATION,
-    DEFAULT_RANSOMWARE_AGENT_CONFIGURATION,
-)
-from monkey_island.cc.models import IslandMode
-from monkey_island.cc.services import IslandModeService
-
-
-@pytest.fixture
-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(IslandMode))
-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",
-    [
-        (IslandMode.UNSET, DEFAULT_AGENT_CONFIGURATION),
-        (IslandMode.ADVANCED, DEFAULT_AGENT_CONFIGURATION),
-        (IslandMode.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