diff --git a/monkey/monkey_island/cc/resources/clear_simulation_data.py b/monkey/monkey_island/cc/resources/clear_simulation_data.py index 342fadd5f..e582dbd11 100644 --- a/monkey/monkey_island/cc/resources/clear_simulation_data.py +++ b/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): - self._repository_service = repository_service + def __init__(self, island_event_queue: IIslandEventQueue): + 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() - return make_response({}, HTTPStatus.OK) + self._island_event_queue.publish(IslandEventTopic.CLEAR_SIMULATION_DATA) + return make_response({}, HTTPStatus.NO_CONTENT) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 0fcf2febc..ab3e856d6 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/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 + ) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_clear_simulation_data.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_clear_simulation_data.py new file mode 100644 index 000000000..3ae081fcb --- /dev/null +++ b/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) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_reset_agent_configuration.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_reset_agent_configuration.py new file mode 100644 index 000000000..c89cb372b --- /dev/null +++ b/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 + ) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index a8369b7da..01d3c5b5f 100644 --- a/vulture_allowlist.py +++ b/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