Merge pull request #2277 from guardicore/2234-refactor-clear-simulation-data

2234 refactor clear simulation data

monkey/monkey_island/cc/resources/clear_simulation_data.py

@@ -2,16 +2,16 @@ from http import HTTPStatus
 
 from flask import make_response
 
+from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
-from monkey_island.cc.services import RepositoryService
 
 
 class ClearSimulationData(AbstractResource):
     urls = ["/api/clear-simulation-data"]
 
-    def __init__(self, repository_service: RepositoryService):
+    def __init__(self, island_event_queue: IIslandEventQueue):
-        self._repository_service = repository_service
+        self._island_event_queue = island_event_queue
 
     @jwt_required
     def post(self):
@@ -19,5 +19,5 @@ class ClearSimulationData(AbstractResource):
         Clear all data collected during the simulation
         """
 
-        self._repository_service.clear_simulation_data()
+        self._island_event_queue.publish(IslandEventTopic.CLEAR_SIMULATION_DATA)
-        return make_response({}, HTTPStatus.OK)
+        return make_response({}, HTTPStatus.NO_CONTENT)

monkey/monkey_island/cc/setup/island_event_handlers.py

@@ -1,11 +1,34 @@
+from functools import partial
+
 from common import DIContainer
-from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
+from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
 from monkey_island.cc.island_event_handlers import reset_agent_configuration
+from monkey_island.cc.repository import ICredentialsRepository
+from monkey_island.cc.services.database import Database
 
 
 def setup_island_event_handlers(container: DIContainer):
-    event_queue = container.resolve(PyPubSubIslandEventQueue)
+    island_event_queue = container.resolve(IIslandEventQueue)
 
-    event_queue.subscribe(
+    _subscribe_reset_agent_configuration_events(island_event_queue, container)
+    _subscribe_clear_simulation_data_events(island_event_queue, container)
+
+
+def _subscribe_reset_agent_configuration_events(
+    island_event_queue: IIslandEventQueue, container: DIContainer
+):
+    island_event_queue.subscribe(
         IslandEventTopic.RESET_AGENT_CONFIGURATION, container.resolve(reset_agent_configuration)
     )
+
+
+def _subscribe_clear_simulation_data_events(
+    island_event_queue: IIslandEventQueue, container: DIContainer
+):
+    legacy_database_reset = partial(Database.reset_db, reset_config=False)
+    island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, legacy_database_reset)
+
+    credentials_repository = container.resolve(ICredentialsRepository)
+    island_event_queue.subscribe(
+        IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials
+    )

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

@@ -0,0 +1,29 @@
+from http import HTTPStatus
+from unittest.mock import MagicMock
+
+import pytest
+from tests.common import StubDIContainer
+
+from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
+from monkey_island.cc.resources import ClearSimulationData
+
+
+@pytest.fixture
+def mock_island_event_queue() -> IIslandEventQueue:
+    return MagicMock(spec=IIslandEventQueue)
+
+
+@pytest.fixture
+def flask_client(build_flask_client, mock_island_event_queue):
+    container = StubDIContainer()
+    container.register_instance(IIslandEventQueue, mock_island_event_queue)
+
+    with build_flask_client(container) as flask_client:
+        yield flask_client
+
+
+def test_clear_simulation_data(flask_client, mock_island_event_queue):
+    resp = flask_client.post(ClearSimulationData.urls[0], follow_redirects=True)
+
+    assert resp.status_code == HTTPStatus.NO_CONTENT
+    mock_island_event_queue.publish.assert_called_once_with(IslandEventTopic.CLEAR_SIMULATION_DATA)

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

@@ -0,0 +1,31 @@
+from http import HTTPStatus
+from unittest.mock import MagicMock
+
+import pytest
+from tests.common import StubDIContainer
+
+from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
+from monkey_island.cc.resources import ResetAgentConfiguration
+
+
+@pytest.fixture
+def mock_island_event_queue() -> IIslandEventQueue:
+    return MagicMock(spec=IIslandEventQueue)
+
+
+@pytest.fixture
+def flask_client(build_flask_client, mock_island_event_queue):
+    container = StubDIContainer()
+    container.register_instance(IIslandEventQueue, mock_island_event_queue)
+
+    with build_flask_client(container) as flask_client:
+        yield flask_client
+
+
+def test_reset_agent_configuration(flask_client, mock_island_event_queue):
+    resp = flask_client.post(ResetAgentConfiguration.urls[0], follow_redirects=True)
+
+    assert resp.status_code == HTTPStatus.OK
+    mock_island_event_queue.publish.assert_called_once_with(
+        IslandEventTopic.RESET_AGENT_CONFIGURATION
+    )

vulture_allowlist.py

@@ -299,6 +299,9 @@ event
 deserialize
 serialized_event
 
+# TODO: Remove once #2277 is merged
+clear_simulation_data
+
 # pydantic base models
 underscore_attrs_are_private
 extra