From bdc5b8fa7b6b054ae1d7db8f47facc4020ebbf06 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Fri, 1 Jul 2022 11:07:49 -0400 Subject: [PATCH] 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)