From 0137c891581374a16f1041a35d952af2421bf8f3 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 08:26:23 -0400 Subject: [PATCH 01/28] Island: Add preliminary Simulation class Includes just the island's mode for now. --- monkey/monkey_island/cc/models/simulation.py | 21 ++++++++++++++++++++ vulture_allowlist.py | 3 +++ 2 files changed, 24 insertions(+) create mode 100644 monkey/monkey_island/cc/models/simulation.py diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py new file mode 100644 index 000000000..9e22c4998 --- /dev/null +++ b/monkey/monkey_island/cc/models/simulation.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from marshmallow import Schema, post_load +from marshmallow_enum import EnumField + +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +@dataclass(frozen=True) +class Simulation: + mode: IslandModeEnum + + +class SimulationSchema(Schema): + mode = EnumField(IslandModeEnum) + + @post_load + def _make_simulation(self, data, **kwargs): + return Simulation(**data) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 126f6c596..0db82d5e1 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -197,6 +197,9 @@ _make_network_scan_configuration # unused method (monkey/common/configuration/a _make_propagation_configuration # unused method (monkey/common/configuration/agent_configuration.py:167) _make_agent_configuration # unused method (monkey/common/configuration/agent_configuration.py:192) +# Models +_make_simulation # unused method (monkey/monkey_island/cc/models/simulation.py:19 + # TODO DELETE AFTER RESOURCE REFACTORING NetworkMap From a84220f34352fb2428799efc043fe082bb60b42a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:13:48 -0400 Subject: [PATCH 02/28] Project: Remove nonexistant _make_agent_configuration from vulture --- vulture_allowlist.py | 1 - 1 file changed, 1 deletion(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 0db82d5e1..22ddb97d5 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -195,7 +195,6 @@ _make_icmp_scan_configuration # unused method (monkey/common/configuration/agen _make_tcp_scan_configuration # unused method (monkey/common/configuration/agent_configuration.py:122) _make_network_scan_configuration # unused method (monkey/common/configuration/agent_configuration.py:110) _make_propagation_configuration # unused method (monkey/common/configuration/agent_configuration.py:167) -_make_agent_configuration # unused method (monkey/common/configuration/agent_configuration.py:192) # Models _make_simulation # unused method (monkey/monkey_island/cc/models/simulation.py:19 From a3f3fd96f72d79b1def91efb056826d93142c1e1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:15:51 -0400 Subject: [PATCH 03/28] Island: Export Simulation from models package --- monkey/monkey_island/cc/models/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index c293ae2e7..2fcf42fd1 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -8,3 +8,4 @@ from .monkey_ttl import MonkeyTtl from .pba_results import PbaResults from monkey_island.cc.models.report.report import Report from .stolen_credentials import StolenCredentials +from .simulation import Simulation, SimulationSchema From 8b10d81d48e09a4d56189adec45075544283a2cf Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:17:19 -0400 Subject: [PATCH 04/28] Island: Use Simulation in ISimulationRepository --- .../monkey_island/cc/repository/i_simulation_repository.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 3a3854fb0..2b41181c8 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,10 +1,10 @@ from abc import ABC +from monkey_island.cc.models import Simulation + class ISimulationRepository(ABC): - # TODO define simulation object. It should contain metadata about simulation, - # like start, end times, mode and last forced stop of all monkeys - def save_simulation(self, simulation: Simulation): # noqa: F821 + def save_simulation(self, simulation: Simulation): pass def get_simulation(self): From a24bf921c1302e9fccc2b9cafb318930099705d7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:19:51 -0400 Subject: [PATCH 05/28] Island: Add {get,set}_mode() to ISimulationRepository --- .../cc/repository/i_simulation_repository.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 2b41181c8..17c5ad9c9 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,6 +1,7 @@ from abc import ABC from monkey_island.cc.models import Simulation +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class ISimulationRepository(ABC): @@ -9,3 +10,22 @@ class ISimulationRepository(ABC): def get_simulation(self): pass + + def get_mode(self) -> IslandModeEnum: + """ + Get's the island's current mode + + :return The island's current mode + :raises RetrievalError: If the mode could not be retrieved + """ + pass + + def set_mode(self, mode: IslandModeEnum): + """ + Set the island's mode + + :param mode: The island's new mode + :raises StorageError: If the mode could not be saved + """ + + pass From 28e75319b02815d66eef3bee92eac40fb01ca67f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:35:38 -0400 Subject: [PATCH 06/28] Island: Make ISimulationRepository methods abstract --- .../monkey_island/cc/repository/i_simulation_repository.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 17c5ad9c9..cf2b06583 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,16 +1,19 @@ -from abc import ABC +from abc import ABC, abstractmethod from monkey_island.cc.models import Simulation from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class ISimulationRepository(ABC): + @abstractmethod def save_simulation(self, simulation: Simulation): pass + @abstractmethod def get_simulation(self): pass + @abstractmethod def get_mode(self) -> IslandModeEnum: """ Get's the island's current mode @@ -20,6 +23,7 @@ class ISimulationRepository(ABC): """ pass + @abstractmethod def set_mode(self, mode: IslandModeEnum): """ Set the island's mode From c221a9810d8105006d05d979d0b09ad8f1250d9a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:36:45 -0400 Subject: [PATCH 07/28] Island: Add missing newline in ISimulationRepository --- monkey/monkey_island/cc/repository/i_simulation_repository.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index cf2b06583..4d7bd23d4 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -21,6 +21,7 @@ class ISimulationRepository(ABC): :return The island's current mode :raises RetrievalError: If the mode could not be retrieved """ + pass @abstractmethod From 55eb9b39a20786e4b49ff53fa48567fd3caf8735 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:39:20 -0400 Subject: [PATCH 08/28] Island: Set default value for Simulation.mode --- monkey/monkey_island/cc/models/simulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py index 9e22c4998..4f64b6a1a 100644 --- a/monkey/monkey_island/cc/models/simulation.py +++ b/monkey/monkey_island/cc/models/simulation.py @@ -10,7 +10,7 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @dataclass(frozen=True) class Simulation: - mode: IslandModeEnum + mode: IslandModeEnum = IslandModeEnum.UNSET class SimulationSchema(Schema): From a18b5ef9b74afe4d80419272f975bb76bf605be9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:41:19 -0400 Subject: [PATCH 09/28] Island: Add missing return typehint to get_simulation() --- monkey/monkey_island/cc/repository/i_simulation_repository.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 4d7bd23d4..b13c46240 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -10,7 +10,7 @@ class ISimulationRepository(ABC): pass @abstractmethod - def get_simulation(self): + def get_simulation(self) -> Simulation: pass @abstractmethod From c58bbb3f8aa9ad5db878739f2f004986788c288b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 09:57:42 -0400 Subject: [PATCH 10/28] Island: Add FileSimulationRepository --- .../monkey_island/cc/repository/__init__.py | 2 + .../repository/file_simulation_repository.py | 41 +++++++++++++++ .../test_file_simulation_repository.py | 52 +++++++++++++++++++ 3 files changed, 95 insertions(+) create mode 100644 monkey/monkey_island/cc/repository/file_simulation_repository.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py diff --git a/monkey/monkey_island/cc/repository/__init__.py b/monkey/monkey_island/cc/repository/__init__.py index 52e6f6873..925a60dd5 100644 --- a/monkey/monkey_island/cc/repository/__init__.py +++ b/monkey/monkey_island/cc/repository/__init__.py @@ -4,3 +4,5 @@ from .i_agent_binary_repository import IAgentBinaryRepository from .agent_binary_repository import AgentBinaryRepository from .i_agent_configuration_repository import IAgentConfigurationRepository from .file_agent_configuration_repository import FileAgentConfigurationRepository +from .i_simulation_repository import ISimulationRepository +from .file_simulation_repository import FileSimulationRepository diff --git a/monkey/monkey_island/cc/repository/file_simulation_repository.py b/monkey/monkey_island/cc/repository/file_simulation_repository.py new file mode 100644 index 000000000..fed7ce0eb --- /dev/null +++ b/monkey/monkey_island/cc/repository/file_simulation_repository.py @@ -0,0 +1,41 @@ +import dataclasses +import io + +from monkey_island.cc import repository +from monkey_island.cc.models import Simulation, SimulationSchema +from monkey_island.cc.repository import IFileRepository, ISimulationRepository, RetrievalError +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + +SIMULATION_STATE_FILE_NAME = "simulation_state.json" + + +class FileSimulationRepository(ISimulationRepository): + def __init__(self, file_repository: IFileRepository): + self._file_repository = file_repository + self._simulation_schema = SimulationSchema() + + def save_simulation(self, simulation: Simulation): + simulation_json = self._simulation_schema.dumps(simulation) + + self._file_repository.save_file( + SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode()) + ) + + def get_simulation(self) -> Simulation: + try: + with self._file_repository.open_file(SIMULATION_STATE_FILE_NAME) as f: + simulation_json = f.read().decode() + + return self._simulation_schema.loads(simulation_json) + except repository.FileNotFoundError: + return Simulation() + except Exception as err: + raise RetrievalError(f"Error retrieving the simulation state: {err}") + + def get_mode(self) -> IslandModeEnum: + return self.get_simulation().mode + + def set_mode(self, mode: IslandModeEnum): + old_simulation = self.get_simulation() + new_simulation = dataclasses.replace(old_simulation, mode=mode) + self.save_simulation(new_simulation) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py new file mode 100644 index 000000000..b5b0f63da --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py @@ -0,0 +1,52 @@ +import pytest +from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository + +from monkey_island.cc.models import Simulation +from monkey_island.cc.repository import FileSimulationRepository, RetrievalError +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +@pytest.fixture +def simulation_repository(): + return FileSimulationRepository(SingleFileRepository()) + + +def test_save_simulation(simulation_repository): + simulation = Simulation(IslandModeEnum.RANSOMWARE) + + old_simulation = simulation_repository.get_simulation() + simulation_repository.save_simulation(simulation) + new_simulation = simulation_repository.get_simulation() + + assert old_simulation != simulation + assert new_simulation == simulation + + +def test_get_default_simulation(simulation_repository): + default_simulation = Simulation() + + assert simulation_repository.get_simulation() == default_simulation + + +def test_set_mode(simulation_repository): + simulation_repository.set_mode(IslandModeEnum.ADVANCED) + + assert simulation_repository.get_mode() == IslandModeEnum.ADVANCED + + +def test_get_mode_default(simulation_repository): + assert simulation_repository.get_mode() == IslandModeEnum.UNSET + + +def test_get_simulation_retrieval_error(): + simulation_repository = FileSimulationRepository(OpenErrorFileRepository()) + + with pytest.raises(RetrievalError): + simulation_repository.get_simulation() + + +def test_get_mode_retrieval_error(): + simulation_repository = FileSimulationRepository(OpenErrorFileRepository()) + + with pytest.raises(RetrievalError): + simulation_repository.get_mode() From d2906ef23bf384207c9ecc83411d4643a73658d7 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 10:00:56 -0400 Subject: [PATCH 11/28] Island: Add docstrings to {get,save}_simulation() --- .../cc/repository/i_simulation_repository.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index b13c46240..880d456fd 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -7,10 +7,22 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class ISimulationRepository(ABC): @abstractmethod def save_simulation(self, simulation: Simulation): + """ + Save the simulation state + + :param simulation: The simulation state + :raises StorageError: If the simulation states could not be saved + """ pass @abstractmethod def get_simulation(self) -> Simulation: + """ + Get the simulation state + + :raises RetrievalError: If the simulation state could not be retrieved + """ + pass @abstractmethod From 350ac6957fe06d2a98ad18003e7ef9b08ff5e097 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 10:59:05 -0400 Subject: [PATCH 12/28] Island: Register FileSimulationRepository with the DIContainer --- monkey/monkey_island/cc/services/initialize.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 922c3654b..04925b9bd 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -8,9 +8,11 @@ from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, FileAgentConfigurationRepository, + FileSimulationRepository, IAgentBinaryRepository, IAgentConfigurationRepository, IFileRepository, + ISimulationRepository, LocalStorageFileRepository, RetrievalError, ) @@ -45,6 +47,7 @@ def initialize_services(data_dir: Path) -> DIContainer: container.register_instance( IAgentConfigurationRepository, container.resolve(FileAgentConfigurationRepository) ) + container.register_instance(ISimulationRepository, container.resolve(FileSimulationRepository)) # This is temporary until we get DI all worked out. PostBreachFilesService.initialize(container.resolve(IFileRepository)) From 942e97cf18012434804001857ad408d533e90be6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:04:22 -0400 Subject: [PATCH 13/28] Island: Decouple ConfigService.reset_config() from island_mode_service --- monkey/monkey_island/cc/services/config.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/monkey/monkey_island/cc/services/config.py b/monkey/monkey_island/cc/services/config.py index 3cb5ce09f..99f5a589c 100644 --- a/monkey/monkey_island/cc/services/config.py +++ b/monkey/monkey_island/cc/services/config.py @@ -23,10 +23,7 @@ from monkey_island.cc.server_utils.encryption import ( encrypt_dict, get_datastore_encryptor, ) -from monkey_island.cc.services.config_manipulator import update_config_per_mode from monkey_island.cc.services.config_schema.config_schema import SCHEMA -from monkey_island.cc.services.mode.island_mode_service import get_mode -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum from monkey_island.cc.services.post_breach_files import PostBreachFilesService logger = logging.getLogger(__name__) @@ -250,14 +247,6 @@ class ConfigService: @staticmethod def reset_config(): PostBreachFilesService.remove_PBA_files() - config = ConfigService.get_default_config(True) - - mode = get_mode() - if mode == IslandModeEnum.UNSET.value: - ConfigService.update_config(config, should_encrypt=False) - else: - update_config_per_mode(mode, config, should_encrypt=False) - logger.info("Monkey config reset was called") @staticmethod From bdc5b8fa7b6b054ae1d7db8f47facc4020ebbf06 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:07:49 -0400 Subject: [PATCH 14/28] Island: Refactor IslandModeService to use ISimulationRepository --- monkey/monkey_island/cc/services/__init__.py | 1 + .../monkey_island/cc/services/initialize.py | 8 +++-- .../cc/services/mode/island_mode_service.py | 32 ++++++++++++------- 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index 66bb1ea65..d655088fc 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -2,3 +2,4 @@ from .authentication.authentication_service import AuthenticationService from .authentication.json_file_user_datastore import JsonFileUserDatastore from .aws import AWSService +from .mode.island_mode_service import IslandModeService diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 04925b9bd..fafd274d3 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -17,7 +17,7 @@ from monkey_island.cc.repository import ( RetrievalError, ) from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH -from monkey_island.cc.services import AWSService +from monkey_island.cc.services import AWSService, IslandModeService from monkey_island.cc.services.post_breach_files import PostBreachFilesService from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService @@ -41,14 +41,16 @@ def initialize_services(data_dir: Path) -> DIContainer: container.register_instance( IFileRepository, LocalStorageFileRepository(data_dir / "runtime_data") ) - container.register_instance(AWSService, container.resolve(AWSService)) container.register_instance(IAgentBinaryRepository, _build_agent_binary_repository()) - container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService)) container.register_instance( IAgentConfigurationRepository, container.resolve(FileAgentConfigurationRepository) ) container.register_instance(ISimulationRepository, container.resolve(FileSimulationRepository)) + container.register_instance(AWSService, container.resolve(AWSService)) + container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService)) + container.register_instance(IslandModeService, container.resolve(IslandModeService)) + # This is temporary until we get DI all worked out. PostBreachFilesService.initialize(container.resolve(IFileRepository)) AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py index 32a3d943d..dd67b6643 100644 --- a/monkey/monkey_island/cc/services/mode/island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -1,17 +1,25 @@ -from monkey_island.cc.models.island_mode_model import IslandMode +from monkey_island.cc.repository import ISimulationRepository from monkey_island.cc.services.mode.mode_enum import IslandModeEnum -def set_mode(mode: IslandModeEnum): - IslandMode.drop_collection() - island_mode_model = IslandMode() - island_mode_model.mode = mode.value - island_mode_model.save() +class IslandModeService: + def __init__(self, simulation_repository: ISimulationRepository): + self._simulation_repository = simulation_repository + def get_mode(self): + """ + Get's the island's current mode -def get_mode() -> str: - if IslandMode.objects: - mode = IslandMode.objects[0].mode - return mode - else: - return IslandModeEnum.UNSET.value + :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: IslandModeEnum): + """ + 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) From 50a982672e6be7cf4e394c7bf2186422fa2cd637 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:17:56 -0400 Subject: [PATCH 15/28] Island: Chande method order in ISimulationRepository --- .../repository/file_simulation_repository.py | 14 ++++++------- .../cc/repository/i_simulation_repository.py | 20 +++++++++---------- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/monkey/monkey_island/cc/repository/file_simulation_repository.py b/monkey/monkey_island/cc/repository/file_simulation_repository.py index fed7ce0eb..021a7083e 100644 --- a/monkey/monkey_island/cc/repository/file_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/file_simulation_repository.py @@ -14,13 +14,6 @@ class FileSimulationRepository(ISimulationRepository): self._file_repository = file_repository self._simulation_schema = SimulationSchema() - def save_simulation(self, simulation: Simulation): - simulation_json = self._simulation_schema.dumps(simulation) - - self._file_repository.save_file( - SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode()) - ) - def get_simulation(self) -> Simulation: try: with self._file_repository.open_file(SIMULATION_STATE_FILE_NAME) as f: @@ -32,6 +25,13 @@ class FileSimulationRepository(ISimulationRepository): except Exception as err: raise RetrievalError(f"Error retrieving the simulation state: {err}") + def save_simulation(self, simulation: Simulation): + simulation_json = self._simulation_schema.dumps(simulation) + + self._file_repository.save_file( + SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode()) + ) + def get_mode(self) -> IslandModeEnum: return self.get_simulation().mode diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 880d456fd..816b691e1 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -5,16 +5,6 @@ from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class ISimulationRepository(ABC): - @abstractmethod - def save_simulation(self, simulation: Simulation): - """ - Save the simulation state - - :param simulation: The simulation state - :raises StorageError: If the simulation states could not be saved - """ - pass - @abstractmethod def get_simulation(self) -> Simulation: """ @@ -25,6 +15,16 @@ class ISimulationRepository(ABC): pass + @abstractmethod + def save_simulation(self, simulation: Simulation): + """ + Save the simulation state + + :param simulation: The simulation state + :raises StorageError: If the simulation states could not be saved + """ + pass + @abstractmethod def get_mode(self) -> IslandModeEnum: """ From 8f7e7f98e7137942aa229fa4606cd5d915e4154a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:31:04 -0400 Subject: [PATCH 16/28] Island: Use new IslandModeService in IslandMode resource --- .../monkey_island/cc/resources/island_mode.py | 21 +++----- monkey/tests/monkey_island/__init__.py | 1 + .../in_memory_simulation_configuration.py | 22 ++++++++ .../cc/resources/test_island_mode.py | 50 +++++++++++-------- 4 files changed, 61 insertions(+), 33 deletions(-) create mode 100644 monkey/tests/monkey_island/in_memory_simulation_configuration.py diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 5d7343459..1a9483eec 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -5,8 +5,7 @@ from flask import make_response, request from monkey_island.cc.resources.AbstractResource import AbstractResource from monkey_island.cc.resources.request_authentication import jwt_required -from monkey_island.cc.services.config_manipulator import update_config_on_mode_set -from monkey_island.cc.services.mode.island_mode_service import get_mode, set_mode +from monkey_island.cc.services import IslandModeService from monkey_island.cc.services.mode.mode_enum import IslandModeEnum logger = logging.getLogger(__name__) @@ -16,20 +15,16 @@ class IslandMode(AbstractResource): # API Spec: Instead of POST, this could just be PATCH urls = ["/api/island-mode"] + def __init__(self, island_mode_service: IslandModeService): + self._island_mode_service = island_mode_service + @jwt_required def post(self): try: body = json.loads(request.data) - mode_str = body.get("mode") + mode = IslandModeEnum(body.get("mode")) - mode = IslandModeEnum(mode_str) - set_mode(mode) - - if not update_config_on_mode_set(mode): - logger.error( - "Could not apply configuration changes per mode. " - "Using default advanced configuration." - ) + self._island_mode_service.set_mode(mode) return make_response({}, 200) except (AttributeError, json.decoder.JSONDecodeError): @@ -39,5 +34,5 @@ class IslandMode(AbstractResource): @jwt_required def get(self): - island_mode = get_mode() - return make_response({"mode": island_mode}, 200) + island_mode = self._island_mode_service.get_mode() + return make_response({"mode": island_mode.value}, 200) diff --git a/monkey/tests/monkey_island/__init__.py b/monkey/tests/monkey_island/__init__.py index 8f3654a6e..935752c37 100644 --- a/monkey/tests/monkey_island/__init__.py +++ b/monkey/tests/monkey_island/__init__.py @@ -2,3 +2,4 @@ from .single_file_repository import SingleFileRepository from .mock_file_repository import MockFileRepository, FILE_CONTENTS, FILE_NAME from .open_error_file_repository import OpenErrorFileRepository from .in_memory_agent_configuration_repository import InMemoryAgentConfigurationRepository +from .in_memory_simulation_configuration import InMemorySimulationRepository diff --git a/monkey/tests/monkey_island/in_memory_simulation_configuration.py b/monkey/tests/monkey_island/in_memory_simulation_configuration.py new file mode 100644 index 000000000..e6e037fcd --- /dev/null +++ b/monkey/tests/monkey_island/in_memory_simulation_configuration.py @@ -0,0 +1,22 @@ +import dataclasses + +from monkey_island.cc.models import Simulation +from monkey_island.cc.repository import ISimulationRepository +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +class InMemorySimulationRepository(ISimulationRepository): + def __init__(self): + self._simulation = Simulation() + + def get_simulation(self) -> Simulation: + return self._simulation + + def save_simulation(self, simulation: Simulation): + self._simulation = simulation + + def get_mode(self) -> IslandModeEnum: + return self._simulation.mode + + def set_mode(self, mode: IslandModeEnum): + self._simulation = dataclasses.replace(self._simulation, mode=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 2e603864d..fec2b6d37 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 @@ -1,28 +1,32 @@ import json +from unittest.mock import MagicMock import pytest -from tests.utils import raise_ +from tests.common import StubDIContainer +from tests.monkey_island import InMemorySimulationRepository -from monkey_island.cc.models.island_mode_model import IslandMode -from monkey_island.cc.resources import island_mode as island_mode_resource +from monkey_island.cc.repository import ISimulationRepository, RetrievalError from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource +from monkey_island.cc.services import IslandModeService from monkey_island.cc.services.mode.mode_enum import IslandModeEnum -@pytest.fixture(scope="function") -def uses_database(): - IslandMode.objects().delete() +@pytest.fixture +def flask_client(build_flask_client): + container = StubDIContainer() + + container.register(ISimulationRepository, InMemorySimulationRepository) + container.register_instance(IslandModeService, container.resolve(IslandModeService)) + + with build_flask_client(container) as flask_client: + yield flask_client @pytest.mark.parametrize( "mode", [IslandModeEnum.RANSOMWARE.value, IslandModeEnum.ADVANCED.value, IslandModeEnum.UNSET.value], ) -def test_island_mode_post(flask_client, mode, monkeypatch): - monkeypatch.setattr( - "monkey_island.cc.resources.island_mode.update_config_on_mode_set", - lambda _: None, - ) +def test_island_mode_post(flask_client, mode): resp = flask_client.post( IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True ) @@ -42,19 +46,25 @@ def test_island_mode_post__invalid_json(flask_client, invalid_json): assert resp.status_code == 400 -def test_island_mode_post__internal_server_error(monkeypatch, flask_client): - monkeypatch.setattr(island_mode_resource, "set_mode", lambda x: raise_(Exception())) +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.post( + IslandModeResource.urls[0], + data=json.dumps({"mode": IslandModeEnum.RANSOMWARE.value}), + follow_redirects=True, + ) - resp = flask_client.post( - IslandModeResource.urls[0], - data=json.dumps({"mode": IslandModeEnum.RANSOMWARE.value}), - follow_redirects=True, - ) assert resp.status_code == 500 @pytest.mark.parametrize("mode", [IslandModeEnum.RANSOMWARE.value, IslandModeEnum.ADVANCED.value]) -def test_island_mode_endpoint(flask_client, uses_database, mode): +def test_island_mode_endpoint(flask_client, mode): flask_client.post( IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True ) @@ -63,7 +73,7 @@ def test_island_mode_endpoint(flask_client, uses_database, mode): assert json.loads(resp.data)["mode"] == mode -def test_island_mode_endpoint__invalid_mode(flask_client, uses_database): +def test_island_mode_endpoint__invalid_mode(flask_client): resp_post = flask_client.post( IslandModeResource.urls[0], data=json.dumps({"mode": "bogus_mode"}), follow_redirects=True ) From 4a91777019c1ccb8481b6cb345713c991fa60334 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:38:20 -0400 Subject: [PATCH 17/28] Island: Add tests for IslandModeService --- .../cc/services/test_island_mode_service.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py 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 new file mode 100644 index 000000000..79074b213 --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py @@ -0,0 +1,16 @@ +import pytest +from tests.monkey_island import InMemorySimulationRepository + +from monkey_island.cc.services import IslandModeService +from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + + +@pytest.fixture +def island_mode_service(): + return IslandModeService(InMemorySimulationRepository()) + + +@pytest.mark.parametrize("mode", list(IslandModeEnum)) +def test_set_mode(island_mode_service, mode): + island_mode_service.set_mode(mode) + assert island_mode_service.get_mode() == mode From 9990cdc6a8b740d57d76db840fb3d4129882cd11 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:40:21 -0400 Subject: [PATCH 18/28] UT: Parametrize test_save_simulation() --- .../cc/repository/test_file_simulation_repository.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py index b5b0f63da..ca4bc6b8c 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py @@ -11,15 +11,12 @@ def simulation_repository(): return FileSimulationRepository(SingleFileRepository()) -def test_save_simulation(simulation_repository): - simulation = Simulation(IslandModeEnum.RANSOMWARE) - - old_simulation = simulation_repository.get_simulation() +@pytest.mark.parametrize("mode", list(IslandModeEnum)) +def test_save_simulation(simulation_repository, mode): + simulation = Simulation(mode) simulation_repository.save_simulation(simulation) - new_simulation = simulation_repository.get_simulation() - assert old_simulation != simulation - assert new_simulation == simulation + assert simulation_repository.get_simulation() == simulation def test_get_default_simulation(simulation_repository): From fc76bb2d83869070766e68bd306d81c25d9d3d38 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:04:48 -0400 Subject: [PATCH 19/28] Island: Set config on mode set in IslandModeService --- monkey/common/configuration/__init__.py | 1 + .../default_agent_configuration.py | 6 ++++ .../monkey_island/cc/services/initialize.py | 11 ++++++- .../cc/services/mode/island_mode_service.py | 25 +++++++++++++-- .../cc/resources/test_island_mode.py | 17 +++++++--- .../cc/services/test_island_mode_service.py | 32 +++++++++++++++++-- 6 files changed, 82 insertions(+), 10 deletions(-) diff --git a/monkey/common/configuration/__init__.py b/monkey/common/configuration/__init__.py index fc1f3c84d..fb52b54e2 100644 --- a/monkey/common/configuration/__init__.py +++ b/monkey/common/configuration/__init__.py @@ -12,4 +12,5 @@ from .agent_sub_configurations import ( ) from .default_agent_configuration import ( DEFAULT_AGENT_CONFIGURATION, + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, ) diff --git a/monkey/common/configuration/default_agent_configuration.py b/monkey/common/configuration/default_agent_configuration.py index 251676017..9045429ca 100644 --- a/monkey/common/configuration/default_agent_configuration.py +++ b/monkey/common/configuration/default_agent_configuration.py @@ -1,3 +1,5 @@ +import dataclasses + from . import AgentConfiguration from .agent_sub_configurations import ( CustomPBAConfiguration, @@ -112,3 +114,7 @@ DEFAULT_AGENT_CONFIGURATION = AgentConfiguration( payloads=PAYLOAD_CONFIGURATION, propagation=PROPAGATION_CONFIGURATION, ) + +DEFAULT_RANSOMWARE_AGENT_CONFIGURATION = dataclasses.replace( + DEFAULT_AGENT_CONFIGURATION, post_breach_actions=[] +) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index fafd274d3..30301b90f 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -3,7 +3,11 @@ from pathlib import Path from common import DIContainer from common.aws import AWSInstance -from common.configuration import DEFAULT_AGENT_CONFIGURATION, AgentConfiguration +from common.configuration import ( + DEFAULT_AGENT_CONFIGURATION, + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, + AgentConfiguration, +) from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.repository import ( AgentBinaryRepository, @@ -36,6 +40,11 @@ def initialize_services(data_dir: Path) -> DIContainer: container.register_convention( AgentConfiguration, "default_agent_configuration", DEFAULT_AGENT_CONFIGURATION ) + container.register_convention( + AgentConfiguration, + "default_ransomware_agent_configuration", + DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, + ) container.register_instance(AWSInstance, AWSInstance()) container.register_instance( diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/mode/island_mode_service.py index dd67b6643..5183bdb8c 100644 --- a/monkey/monkey_island/cc/services/mode/island_mode_service.py +++ b/monkey/monkey_island/cc/services/mode/island_mode_service.py @@ -1,10 +1,20 @@ -from monkey_island.cc.repository import ISimulationRepository +from common.configuration import AgentConfiguration +from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class IslandModeService: - def __init__(self, simulation_repository: ISimulationRepository): + 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): """ @@ -23,3 +33,14 @@ class IslandModeService: :raises StorageError: If the mode could not be saved """ self._simulation_repository.set_mode(mode) + self._set_configuration(mode) + + def _set_configuration(self, mode: IslandModeEnum): + if mode == IslandModeEnum.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/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 fec2b6d37..ae3905923 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,18 +5,27 @@ import pytest from tests.common import StubDIContainer from tests.monkey_island import InMemorySimulationRepository -from monkey_island.cc.repository import ISimulationRepository, RetrievalError +from monkey_island.cc.repository import RetrievalError from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource from monkey_island.cc.services import IslandModeService from monkey_island.cc.services.mode.mode_enum import IslandModeEnum +class MockIslandModeService(IslandModeService): + def __init__(self): + self._simulation_repository = InMemorySimulationRepository() + + def get_mode(self) -> IslandModeEnum: + return self._simulation_repository.get_mode() + + def set_mode(self, mode: IslandModeEnum): + self._simulation_repository.set_mode(mode) + + @pytest.fixture def flask_client(build_flask_client): container = StubDIContainer() - - container.register(ISimulationRepository, InMemorySimulationRepository) - container.register_instance(IslandModeService, container.resolve(IslandModeService)) + container.register_instance(IslandModeService, MockIslandModeService()) with build_flask_client(container) as flask_client: yield flask_client 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 index 79074b213..4d5f5f4bc 100644 --- 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 @@ -1,16 +1,42 @@ import pytest -from tests.monkey_island import InMemorySimulationRepository +from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository +from common.configuration import DEFAULT_AGENT_CONFIGURATION, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION from monkey_island.cc.services import IslandModeService from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @pytest.fixture -def island_mode_service(): - return IslandModeService(InMemorySimulationRepository()) +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(IslandModeEnum)) 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", + [ + (IslandModeEnum.UNSET, DEFAULT_AGENT_CONFIGURATION), + (IslandModeEnum.ADVANCED, DEFAULT_AGENT_CONFIGURATION), + (IslandModeEnum.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 From c1994f5d607319054846a2677105b243d8fe6609 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:09:29 -0400 Subject: [PATCH 20/28] Island: Extract registration methods from initialize_services() --- .../monkey_island/cc/services/initialize.py | 33 ++++++++++++------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 30301b90f..28a88558c 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -35,7 +35,20 @@ AGENT_BINARIES_PATH = Path(MONKEY_ISLAND_ABS_PATH) / "cc" / "binaries" def initialize_services(data_dir: Path) -> DIContainer: container = DIContainer() + _register_conventions(container, data_dir) + container.register_instance(AWSInstance, AWSInstance()) + _register_repositories(container, data_dir) + _register_services(container) + # This is temporary until we get DI all worked out. + PostBreachFilesService.initialize(container.resolve(IFileRepository)) + AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) + ReportService.initialize(container.resolve(AWSService)) + + return container + + +def _register_conventions(container: DIContainer, data_dir: Path): container.register_convention(Path, "data_dir", data_dir) container.register_convention( AgentConfiguration, "default_agent_configuration", DEFAULT_AGENT_CONFIGURATION @@ -45,8 +58,9 @@ def initialize_services(data_dir: Path) -> DIContainer: "default_ransomware_agent_configuration", DEFAULT_RANSOMWARE_AGENT_CONFIGURATION, ) - container.register_instance(AWSInstance, AWSInstance()) + +def _register_repositories(container: DIContainer, data_dir: Path): container.register_instance( IFileRepository, LocalStorageFileRepository(data_dir / "runtime_data") ) @@ -56,17 +70,6 @@ def initialize_services(data_dir: Path) -> DIContainer: ) container.register_instance(ISimulationRepository, container.resolve(FileSimulationRepository)) - container.register_instance(AWSService, container.resolve(AWSService)) - container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService)) - container.register_instance(IslandModeService, container.resolve(IslandModeService)) - - # This is temporary until we get DI all worked out. - PostBreachFilesService.initialize(container.resolve(IFileRepository)) - AuthenticationService.initialize(data_dir, JsonFileUserDatastore(data_dir)) - ReportService.initialize(container.resolve(AWSService)) - - return container - def _build_agent_binary_repository(): file_repository = LocalStorageFileRepository(AGENT_BINARIES_PATH) @@ -99,3 +102,9 @@ def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository): for os, binary_sha256_hash in agent_hashes.items(): logger.info(f"{os} agent: SHA-256 hash: {binary_sha256_hash}") + + +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)) From 6459da1d15da70975f7f407903f8e2cbd89e550c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:14:52 -0400 Subject: [PATCH 21/28] Island: Remove disused config_manipulator "service" --- .../cc/services/config_manipulator.py | 31 ------------------ .../cc/services/test_config_manipulator.py | 32 ------------------- 2 files changed, 63 deletions(-) delete mode 100644 monkey/monkey_island/cc/services/config_manipulator.py delete mode 100644 monkey/tests/unit_tests/monkey_island/cc/services/test_config_manipulator.py diff --git a/monkey/monkey_island/cc/services/config_manipulator.py b/monkey/monkey_island/cc/services/config_manipulator.py deleted file mode 100644 index 41f555be4..000000000 --- a/monkey/monkey_island/cc/services/config_manipulator.py +++ /dev/null @@ -1,31 +0,0 @@ -from typing import Dict - -import dpath.util - -import monkey_island.cc.services.config as config_service -from monkey_island.cc.services.config_manipulators import MANIPULATOR_PER_MODE -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum - - -def update_config_on_mode_set(mode: IslandModeEnum) -> bool: - config = config_service.ConfigService.get_config() - return update_config_per_mode(mode.value, config, True) - - -def update_config_per_mode(mode: str, config: Dict, should_encrypt: bool) -> bool: - config = _set_default_config_values_per_mode(mode, config) - return config_service.ConfigService.update_config( - config_json=config, should_encrypt=should_encrypt - ) - - -def _set_default_config_values_per_mode(mode: str, config: Dict) -> Dict: - config_manipulator = MANIPULATOR_PER_MODE[mode] - config = _apply_config_manipulator(config, config_manipulator) - return config - - -def _apply_config_manipulator(config: Dict, config_manipulator: Dict): - for path, value in config_manipulator.items(): - dpath.util.set(config, path, value, ".") - return config diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_config_manipulator.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_config_manipulator.py deleted file mode 100644 index 403b5aee3..000000000 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_config_manipulator.py +++ /dev/null @@ -1,32 +0,0 @@ -import pytest - -from monkey_island.cc.services.config_manipulator import update_config_on_mode_set -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum - - -@pytest.mark.slow -@pytest.mark.usefixtures("uses_encryptor") -def test_update_config_on_mode_set_advanced(config, monkeypatch): - monkeypatch.setattr("monkey_island.cc.services.config.ConfigService.get_config", lambda: config) - monkeypatch.setattr( - "monkey_island.cc.services.config.ConfigService.update_config", - lambda config_json, should_encrypt: config_json, - ) - - mode = IslandModeEnum.ADVANCED - manipulated_config = update_config_on_mode_set(mode) - assert manipulated_config == config - - -@pytest.mark.slow -@pytest.mark.usefixtures("uses_encryptor") -def test_update_config_on_mode_set_ransomware(config, monkeypatch): - monkeypatch.setattr("monkey_island.cc.services.config.ConfigService.get_config", lambda: config) - monkeypatch.setattr( - "monkey_island.cc.services.config.ConfigService.update_config", - lambda config_json, should_encrypt: config_json, - ) - - mode = IslandModeEnum.RANSOMWARE - manipulated_config = update_config_on_mode_set(mode) - assert manipulated_config["monkey"]["post_breach"]["post_breach_actions"] == [] From 77baf08ea311416221c6897f5d8d81630ecf9110 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:15:42 -0400 Subject: [PATCH 22/28] Island: Remove disused MANIPULATOR_PER_MODE "service" --- monkey/monkey_island/cc/services/config_manipulators.py | 7 ------- 1 file changed, 7 deletions(-) delete mode 100644 monkey/monkey_island/cc/services/config_manipulators.py diff --git a/monkey/monkey_island/cc/services/config_manipulators.py b/monkey/monkey_island/cc/services/config_manipulators.py deleted file mode 100644 index 947a32971..000000000 --- a/monkey/monkey_island/cc/services/config_manipulators.py +++ /dev/null @@ -1,7 +0,0 @@ -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum - -MANIPULATOR_PER_MODE = { - IslandModeEnum.ADVANCED.value: {}, - IslandModeEnum.RANSOMWARE.value: {"monkey.post_breach.post_breach_actions": []}, - IslandModeEnum.UNSET.value: {}, -} From 7ec558a5f448cddf4ab8dc349ee0479fc570b035 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:17:10 -0400 Subject: [PATCH 23/28] Island: Move IslandModeService out of unnecessary subpackage --- monkey/monkey_island/cc/services/__init__.py | 2 +- .../monkey_island/cc/services/{mode => }/island_mode_service.py | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename monkey/monkey_island/cc/services/{mode => }/island_mode_service.py (100%) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index d655088fc..bfecc7588 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -2,4 +2,4 @@ from .authentication.authentication_service import AuthenticationService from .authentication.json_file_user_datastore import JsonFileUserDatastore from .aws import AWSService -from .mode.island_mode_service import IslandModeService +from .island_mode_service import IslandModeService diff --git a/monkey/monkey_island/cc/services/mode/island_mode_service.py b/monkey/monkey_island/cc/services/island_mode_service.py similarity index 100% rename from monkey/monkey_island/cc/services/mode/island_mode_service.py rename to monkey/monkey_island/cc/services/island_mode_service.py From 63a69b047d3d614aad4d0688c021c14cecc7c69b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:20:57 -0400 Subject: [PATCH 24/28] Island: Move IslandModeEnum to monkey_island.cc.models --- monkey/monkey_island/cc/models/__init__.py | 2 +- monkey/monkey_island/cc/models/simulation.py | 7 ++++++- .../cc/repository/file_simulation_repository.py | 3 +-- .../monkey_island/cc/repository/i_simulation_repository.py | 3 +-- monkey/monkey_island/cc/resources/island_mode.py | 2 +- monkey/monkey_island/cc/services/island_mode_service.py | 2 +- .../monkey_island/in_memory_simulation_configuration.py | 3 +-- .../cc/repository/test_file_simulation_repository.py | 3 +-- .../monkey_island/cc/resources/test_island_mode.py | 2 +- .../monkey_island/cc/services/test_island_mode_service.py | 2 +- 10 files changed, 15 insertions(+), 14 deletions(-) diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 2fcf42fd1..7f7cde2d7 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -8,4 +8,4 @@ from .monkey_ttl import MonkeyTtl from .pba_results import PbaResults from monkey_island.cc.models.report.report import Report from .stolen_credentials import StolenCredentials -from .simulation import Simulation, SimulationSchema +from .simulation import Simulation, SimulationSchema, IslandModeEnum diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py index 4f64b6a1a..693287160 100644 --- a/monkey/monkey_island/cc/models/simulation.py +++ b/monkey/monkey_island/cc/models/simulation.py @@ -1,11 +1,16 @@ from __future__ import annotations from dataclasses import dataclass +from enum import Enum from marshmallow import Schema, post_load from marshmallow_enum import EnumField -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum + +class IslandModeEnum(Enum): + UNSET = "unset" + RANSOMWARE = "ransomware" + ADVANCED = "advanced" @dataclass(frozen=True) diff --git a/monkey/monkey_island/cc/repository/file_simulation_repository.py b/monkey/monkey_island/cc/repository/file_simulation_repository.py index 021a7083e..66db53a7f 100644 --- a/monkey/monkey_island/cc/repository/file_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/file_simulation_repository.py @@ -2,9 +2,8 @@ import dataclasses import io from monkey_island.cc import repository -from monkey_island.cc.models import Simulation, SimulationSchema +from monkey_island.cc.models import IslandModeEnum, Simulation, SimulationSchema from monkey_island.cc.repository import IFileRepository, ISimulationRepository, RetrievalError -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum SIMULATION_STATE_FILE_NAME = "simulation_state.json" diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 816b691e1..80dda8076 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,7 +1,6 @@ from abc import ABC, abstractmethod -from monkey_island.cc.models import Simulation -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum +from monkey_island.cc.models import IslandModeEnum, Simulation class ISimulationRepository(ABC): diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 1a9483eec..43ca46d2d 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -3,10 +3,10 @@ import logging from flask import make_response, request +from monkey_island.cc.models import IslandModeEnum 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 -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum logger = logging.getLogger(__name__) diff --git a/monkey/monkey_island/cc/services/island_mode_service.py b/monkey/monkey_island/cc/services/island_mode_service.py index 5183bdb8c..7bb08aa60 100644 --- a/monkey/monkey_island/cc/services/island_mode_service.py +++ b/monkey/monkey_island/cc/services/island_mode_service.py @@ -1,6 +1,6 @@ from common.configuration import AgentConfiguration +from monkey_island.cc.models import IslandModeEnum from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class IslandModeService: diff --git a/monkey/tests/monkey_island/in_memory_simulation_configuration.py b/monkey/tests/monkey_island/in_memory_simulation_configuration.py index e6e037fcd..7044889d2 100644 --- a/monkey/tests/monkey_island/in_memory_simulation_configuration.py +++ b/monkey/tests/monkey_island/in_memory_simulation_configuration.py @@ -1,8 +1,7 @@ import dataclasses -from monkey_island.cc.models import Simulation +from monkey_island.cc.models import IslandModeEnum, Simulation from monkey_island.cc.repository import ISimulationRepository -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class InMemorySimulationRepository(ISimulationRepository): diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py index ca4bc6b8c..3ed23b83a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py @@ -1,9 +1,8 @@ import pytest from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from monkey_island.cc.models import Simulation +from monkey_island.cc.models import IslandModeEnum, Simulation from monkey_island.cc.repository import FileSimulationRepository, RetrievalError -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @pytest.fixture 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 ae3905923..d875fddb8 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,10 +5,10 @@ import pytest from tests.common import StubDIContainer from tests.monkey_island import InMemorySimulationRepository +from monkey_island.cc.models import IslandModeEnum from monkey_island.cc.repository import RetrievalError from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource from monkey_island.cc.services import IslandModeService -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum class MockIslandModeService(IslandModeService): 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 index 4d5f5f4bc..9e1ce30bf 100644 --- 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 @@ -2,8 +2,8 @@ import pytest from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository from common.configuration import DEFAULT_AGENT_CONFIGURATION, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION +from monkey_island.cc.models import IslandModeEnum from monkey_island.cc.services import IslandModeService -from monkey_island.cc.services.mode.mode_enum import IslandModeEnum @pytest.fixture From d88fc86d9018c2506aefd9722a6da181d37433e9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:39:24 -0400 Subject: [PATCH 25/28] Island: Remove disused mode subpackage --- monkey/monkey_island/cc/services/mode/__init__.py | 0 monkey/monkey_island/cc/services/mode/mode_enum.py | 7 ------- 2 files changed, 7 deletions(-) delete mode 100644 monkey/monkey_island/cc/services/mode/__init__.py delete mode 100644 monkey/monkey_island/cc/services/mode/mode_enum.py diff --git a/monkey/monkey_island/cc/services/mode/__init__.py b/monkey/monkey_island/cc/services/mode/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/monkey/monkey_island/cc/services/mode/mode_enum.py b/monkey/monkey_island/cc/services/mode/mode_enum.py deleted file mode 100644 index d6258bbff..000000000 --- a/monkey/monkey_island/cc/services/mode/mode_enum.py +++ /dev/null @@ -1,7 +0,0 @@ -from enum import Enum - - -class IslandModeEnum(Enum): - UNSET = "unset" - RANSOMWARE = "ransomware" - ADVANCED = "advanced" From f78fa73563c181f9268c09d2d6598d2e10c9fe3b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:43:13 -0400 Subject: [PATCH 26/28] Island: Rename IslandModeEnum -> IslandMode --- monkey/monkey_island/cc/models/__init__.py | 2 +- monkey/monkey_island/cc/models/simulation.py | 6 +++--- .../cc/repository/file_simulation_repository.py | 6 +++--- .../cc/repository/i_simulation_repository.py | 6 +++--- monkey/monkey_island/cc/resources/island_mode.py | 2 +- .../cc/services/island_mode_service.py | 8 ++++---- .../in_memory_simulation_configuration.py | 6 +++--- .../repository/test_file_simulation_repository.py | 10 +++++----- .../monkey_island/cc/resources/test_island_mode.py | 14 +++++++------- .../cc/services/test_island_mode_service.py | 10 +++++----- 10 files changed, 35 insertions(+), 35 deletions(-) diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index 7f7cde2d7..bf2addf63 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/monkey/monkey_island/cc/models/__init__.py @@ -8,4 +8,4 @@ from .monkey_ttl import MonkeyTtl from .pba_results import PbaResults from monkey_island.cc.models.report.report import Report from .stolen_credentials import StolenCredentials -from .simulation import Simulation, SimulationSchema, IslandModeEnum +from .simulation import Simulation, SimulationSchema, IslandMode diff --git a/monkey/monkey_island/cc/models/simulation.py b/monkey/monkey_island/cc/models/simulation.py index 693287160..88b3ca25d 100644 --- a/monkey/monkey_island/cc/models/simulation.py +++ b/monkey/monkey_island/cc/models/simulation.py @@ -7,7 +7,7 @@ from marshmallow import Schema, post_load from marshmallow_enum import EnumField -class IslandModeEnum(Enum): +class IslandMode(Enum): UNSET = "unset" RANSOMWARE = "ransomware" ADVANCED = "advanced" @@ -15,11 +15,11 @@ class IslandModeEnum(Enum): @dataclass(frozen=True) class Simulation: - mode: IslandModeEnum = IslandModeEnum.UNSET + mode: IslandMode = IslandMode.UNSET class SimulationSchema(Schema): - mode = EnumField(IslandModeEnum) + mode = EnumField(IslandMode) @post_load def _make_simulation(self, data, **kwargs): diff --git a/monkey/monkey_island/cc/repository/file_simulation_repository.py b/monkey/monkey_island/cc/repository/file_simulation_repository.py index 66db53a7f..0b4843fb6 100644 --- a/monkey/monkey_island/cc/repository/file_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/file_simulation_repository.py @@ -2,7 +2,7 @@ import dataclasses import io from monkey_island.cc import repository -from monkey_island.cc.models import IslandModeEnum, Simulation, SimulationSchema +from monkey_island.cc.models import IslandMode, Simulation, SimulationSchema from monkey_island.cc.repository import IFileRepository, ISimulationRepository, RetrievalError SIMULATION_STATE_FILE_NAME = "simulation_state.json" @@ -31,10 +31,10 @@ class FileSimulationRepository(ISimulationRepository): SIMULATION_STATE_FILE_NAME, io.BytesIO(simulation_json.encode()) ) - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: return self.get_simulation().mode - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): old_simulation = self.get_simulation() new_simulation = dataclasses.replace(old_simulation, mode=mode) self.save_simulation(new_simulation) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index 80dda8076..e6230c724 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod -from monkey_island.cc.models import IslandModeEnum, Simulation +from monkey_island.cc.models import IslandMode, Simulation class ISimulationRepository(ABC): @@ -25,7 +25,7 @@ class ISimulationRepository(ABC): pass @abstractmethod - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: """ Get's the island's current mode @@ -36,7 +36,7 @@ class ISimulationRepository(ABC): pass @abstractmethod - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): """ Set the island's mode diff --git a/monkey/monkey_island/cc/resources/island_mode.py b/monkey/monkey_island/cc/resources/island_mode.py index 43ca46d2d..3273f150e 100644 --- a/monkey/monkey_island/cc/resources/island_mode.py +++ b/monkey/monkey_island/cc/resources/island_mode.py @@ -3,7 +3,7 @@ import logging from flask import make_response, request -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode as IslandModeEnum 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 diff --git a/monkey/monkey_island/cc/services/island_mode_service.py b/monkey/monkey_island/cc/services/island_mode_service.py index 7bb08aa60..2d65d0067 100644 --- a/monkey/monkey_island/cc/services/island_mode_service.py +++ b/monkey/monkey_island/cc/services/island_mode_service.py @@ -1,5 +1,5 @@ from common.configuration import AgentConfiguration -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository @@ -25,7 +25,7 @@ class IslandModeService: """ return self._simulation_repository.get_mode() - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): """ Set the island's mode @@ -35,8 +35,8 @@ class IslandModeService: self._simulation_repository.set_mode(mode) self._set_configuration(mode) - def _set_configuration(self, mode: IslandModeEnum): - if mode == IslandModeEnum.RANSOMWARE: + def _set_configuration(self, mode: IslandMode): + if mode == IslandMode.RANSOMWARE: self._agent_configuration_repository.store_configuration( self._default_ransomware_agent_configuration ) diff --git a/monkey/tests/monkey_island/in_memory_simulation_configuration.py b/monkey/tests/monkey_island/in_memory_simulation_configuration.py index 7044889d2..b2b90fb40 100644 --- a/monkey/tests/monkey_island/in_memory_simulation_configuration.py +++ b/monkey/tests/monkey_island/in_memory_simulation_configuration.py @@ -1,6 +1,6 @@ import dataclasses -from monkey_island.cc.models import IslandModeEnum, Simulation +from monkey_island.cc.models import IslandMode, Simulation from monkey_island.cc.repository import ISimulationRepository @@ -14,8 +14,8 @@ class InMemorySimulationRepository(ISimulationRepository): def save_simulation(self, simulation: Simulation): self._simulation = simulation - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: return self._simulation.mode - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): self._simulation = dataclasses.replace(self._simulation, mode=mode) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py index 3ed23b83a..d01775b98 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_file_simulation_repository.py @@ -1,7 +1,7 @@ import pytest from tests.monkey_island import OpenErrorFileRepository, SingleFileRepository -from monkey_island.cc.models import IslandModeEnum, Simulation +from monkey_island.cc.models import IslandMode, Simulation from monkey_island.cc.repository import FileSimulationRepository, RetrievalError @@ -10,7 +10,7 @@ def simulation_repository(): return FileSimulationRepository(SingleFileRepository()) -@pytest.mark.parametrize("mode", list(IslandModeEnum)) +@pytest.mark.parametrize("mode", list(IslandMode)) def test_save_simulation(simulation_repository, mode): simulation = Simulation(mode) simulation_repository.save_simulation(simulation) @@ -25,13 +25,13 @@ def test_get_default_simulation(simulation_repository): def test_set_mode(simulation_repository): - simulation_repository.set_mode(IslandModeEnum.ADVANCED) + simulation_repository.set_mode(IslandMode.ADVANCED) - assert simulation_repository.get_mode() == IslandModeEnum.ADVANCED + assert simulation_repository.get_mode() == IslandMode.ADVANCED def test_get_mode_default(simulation_repository): - assert simulation_repository.get_mode() == IslandModeEnum.UNSET + assert simulation_repository.get_mode() == IslandMode.UNSET def test_get_simulation_retrieval_error(): 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 d875fddb8..780553eaa 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,7 +5,7 @@ import pytest from tests.common import StubDIContainer from tests.monkey_island import InMemorySimulationRepository -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode from monkey_island.cc.repository import RetrievalError from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource from monkey_island.cc.services import IslandModeService @@ -15,10 +15,10 @@ class MockIslandModeService(IslandModeService): def __init__(self): self._simulation_repository = InMemorySimulationRepository() - def get_mode(self) -> IslandModeEnum: + def get_mode(self) -> IslandMode: return self._simulation_repository.get_mode() - def set_mode(self, mode: IslandModeEnum): + def set_mode(self, mode: IslandMode): self._simulation_repository.set_mode(mode) @@ -33,7 +33,7 @@ def flask_client(build_flask_client): @pytest.mark.parametrize( "mode", - [IslandModeEnum.RANSOMWARE.value, IslandModeEnum.ADVANCED.value, IslandModeEnum.UNSET.value], + [IslandMode.RANSOMWARE.value, IslandMode.ADVANCED.value, IslandMode.UNSET.value], ) def test_island_mode_post(flask_client, mode): resp = flask_client.post( @@ -65,14 +65,14 @@ def test_island_mode_post__internal_server_error(build_flask_client): with build_flask_client(container) as flask_client: resp = flask_client.post( IslandModeResource.urls[0], - data=json.dumps({"mode": IslandModeEnum.RANSOMWARE.value}), + data=json.dumps({"mode": IslandMode.RANSOMWARE.value}), follow_redirects=True, ) assert resp.status_code == 500 -@pytest.mark.parametrize("mode", [IslandModeEnum.RANSOMWARE.value, IslandModeEnum.ADVANCED.value]) +@pytest.mark.parametrize("mode", [IslandMode.RANSOMWARE.value, IslandMode.ADVANCED.value]) def test_island_mode_endpoint(flask_client, mode): flask_client.post( IslandModeResource.urls[0], data=json.dumps({"mode": mode}), follow_redirects=True @@ -88,4 +88,4 @@ def test_island_mode_endpoint__invalid_mode(flask_client): ) resp_get = flask_client.get(IslandModeResource.urls[0], follow_redirects=True) assert resp_post.status_code == 422 - assert json.loads(resp_get.data)["mode"] == IslandModeEnum.UNSET.value + assert json.loads(resp_get.data)["mode"] == IslandMode.UNSET.value 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 index 9e1ce30bf..feb125b49 100644 --- 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 @@ -2,7 +2,7 @@ import pytest from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository from common.configuration import DEFAULT_AGENT_CONFIGURATION, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION -from monkey_island.cc.models import IslandModeEnum +from monkey_island.cc.models import IslandMode from monkey_island.cc.services import IslandModeService @@ -21,7 +21,7 @@ def island_mode_service(agent_configuration_repository): ) -@pytest.mark.parametrize("mode", list(IslandModeEnum)) +@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 @@ -30,9 +30,9 @@ def test_set_mode(island_mode_service, mode): @pytest.mark.parametrize( "mode, expected_config", [ - (IslandModeEnum.UNSET, DEFAULT_AGENT_CONFIGURATION), - (IslandModeEnum.ADVANCED, DEFAULT_AGENT_CONFIGURATION), - (IslandModeEnum.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION), + (IslandMode.UNSET, DEFAULT_AGENT_CONFIGURATION), + (IslandMode.ADVANCED, DEFAULT_AGENT_CONFIGURATION), + (IslandMode.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION), ], ) def test_set_mode_sets_config( From 410401a3197833b638fc2f38e670910fca5355e0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 12:44:35 -0400 Subject: [PATCH 27/28] Island; Remove disused IslandMode mongo document --- monkey/monkey_island/cc/models/island_mode_model.py | 7 ------- monkey/monkey_island/cc/services/database.py | 3 +-- 2 files changed, 1 insertion(+), 9 deletions(-) delete mode 100644 monkey/monkey_island/cc/models/island_mode_model.py diff --git a/monkey/monkey_island/cc/models/island_mode_model.py b/monkey/monkey_island/cc/models/island_mode_model.py deleted file mode 100644 index 8d6aab74a..000000000 --- a/monkey/monkey_island/cc/models/island_mode_model.py +++ /dev/null @@ -1,7 +0,0 @@ -from mongoengine import Document, StringField - - -class IslandMode(Document): - COLLECTION_NAME = "island_mode" - - mode = StringField() diff --git a/monkey/monkey_island/cc/services/database.py b/monkey/monkey_island/cc/services/database.py index 46b5e0ffd..3dd512026 100644 --- a/monkey/monkey_island/cc/services/database.py +++ b/monkey/monkey_island/cc/services/database.py @@ -6,7 +6,6 @@ from monkey_island.cc.database import mongo from monkey_island.cc.models import Config from monkey_island.cc.models.agent_controls import AgentControls from monkey_island.cc.models.attack.attack_mitigations import AttackMitigations -from monkey_island.cc.models.island_mode_model import IslandMode from monkey_island.cc.services.config import ConfigService logger = logging.getLogger(__name__) @@ -33,7 +32,7 @@ class Database(object): @staticmethod def _should_drop(collection: str, drop_config: bool) -> bool: if not drop_config: - if collection == IslandMode.COLLECTION_NAME or collection == Config.COLLECTION_NAME: + if collection == Config.COLLECTION_NAME: return False return ( not collection.startswith("system.") From 1421b5624ca7bd2263bb63afb6de16c9d1bc983b Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 16:31:55 -0400 Subject: [PATCH 28/28] Island: Add missing newline in ISimulationRepository --- monkey/monkey_island/cc/repository/i_simulation_repository.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/monkey_island/cc/repository/i_simulation_repository.py b/monkey/monkey_island/cc/repository/i_simulation_repository.py index e6230c724..94bef50c4 100644 --- a/monkey/monkey_island/cc/repository/i_simulation_repository.py +++ b/monkey/monkey_island/cc/repository/i_simulation_repository.py @@ -22,6 +22,7 @@ class ISimulationRepository(ABC): :param simulation: The simulation state :raises StorageError: If the simulation states could not be saved """ + pass @abstractmethod