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

2234 refactor clear simulation data
This commit is contained in:
Mike Salvatore 2022-09-13 07:25:49 -04:00 committed by GitHub
commit b226c84b6c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 94 additions and 8 deletions

View File

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

View File

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

View File

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

View File

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

View File

@ -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