Island: Refactor IslandModeService to use ISimulationRepository

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

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

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)