diff --git a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py index d2398df86..cc81f8e26 100644 --- a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py +++ b/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): diff --git a/monkey/monkey_island/cc/island_event_handlers/__init__.py b/monkey/monkey_island/cc/island_event_handlers/__init__.py index 0c977976e..11343c2fa 100644 --- a/monkey/monkey_island/cc/island_event_handlers/__init__.py +++ b/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 diff --git a/monkey/monkey_island/cc/island_event_handlers/set_agent_configuration_per_island_mode.py b/monkey/monkey_island/cc/island_event_handlers/set_agent_configuration_per_island_mode.py new file mode 100644 index 000000000..8385d0480 --- /dev/null +++ b/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 + ) diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 84d2858c5..f07e791aa 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/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): - self._island_mode_service = island_mode_service + def __init__( + 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_mode_service.set_mode(mode) - + self._island_event_queue.publish(topic=IslandEventTopic.SET_ISLAND_MODE, 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 diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 14c5664e9..d75734d59 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/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 diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 9e6425215..830e010d7 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/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)) diff --git a/monkey/monkey_island/cc/services/island_mode_service.py b/monkey/monkey_island/cc/services/island_mode_service.py deleted file mode 100644 index bc4869b9a..000000000 --- a/monkey/monkey_island/cc/services/island_mode_service.py +++ /dev/null @@ -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 - ) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index a69cf1731..ee37568f1 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/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( - IslandEventTopic.RESET_AGENT_CONFIGURATION, container.resolve(reset_agent_configuration) - ) + topic = IslandEventTopic.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( - IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials - ) + island_event_queue.subscribe(topic, credentials_repository.remove_stolen_credentials) - island_event_queue.subscribe( - IslandEventTopic.CLEAR_SIMULATION_DATA, container.resolve(reset_machine_repository) - ) + island_event_queue.subscribe(topic, 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) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py index 322f62502..ad099cd9c 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py +++ b/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): - container = StubDIContainer() - container.register_instance(IslandModeService, MockIslandModeService()) +def flask_client_builder(build_flask_client): + def inner(side_effect=None): + container = StubDIContainer() - with build_flask_client(container) as flask_client: - yield flask_client + 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: + 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): - mock_island_mode_service = MagicMock(spec=IslandModeService) - 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, - ) - +def test_island_mode_post__internal_server_error(flask_client__internal_server_error): + resp = flask_client__internal_server_error.put( + IslandModeResource.urls[0], + json=IslandMode.RANSOMWARE.value, + follow_redirects=True, + ) assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py deleted file mode 100644 index 1d12efb62..000000000 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py +++ /dev/null @@ -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