Merge pull request #2086 from guardicore/repository-service-clear-simulation-data

Repository service clear simulation data

monkey/monkey_island/cc/resources/clear_simulation_data.py

@@ -2,24 +2,22 @@ from http import HTTPStatus
 
 from flask import make_response
 
-from monkey_island.cc.repository.i_credentials_repository import ICredentialsRepository
 from monkey_island.cc.resources.AbstractResource import AbstractResource
 from monkey_island.cc.resources.request_authentication import jwt_required
-from monkey_island.cc.services.database import Database
+from monkey_island.cc.services import RepositoryService
 
 
 class ClearSimulationData(AbstractResource):
     urls = ["/api/clear-simulation-data"]
 
-    def __init__(self, credentials_repository: ICredentialsRepository):
-        self._credentials_repository = credentials_repository
+    def __init__(self, repository_service: RepositoryService):
+        self._repository_service = repository_service
 
     @jwt_required
     def post(self):
         """
         Clear all data collected during the simulation
         """
-        Database.reset_db(reset_config=False)
-        self._credentials_repository.remove_stolen_credentials()
 
+        self._repository_service.clear_simulation_data()
         return make_response({}, HTTPStatus.OK)

monkey/monkey_island/cc/services/repository_service.py

@@ -1,4 +1,9 @@
-from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
+from monkey_island.cc.repository import (
+    IAgentConfigurationRepository,
+    ICredentialsRepository,
+    IFileRepository,
+)
+from monkey_island.cc.services.database import Database
 
 
 class RepositoryService:
@@ -6,9 +11,11 @@ class RepositoryService:
         self,
         agent_configuration_repository: IAgentConfigurationRepository,
         file_repository: IFileRepository,
+        credentials_repository: ICredentialsRepository,
     ):
         self._agent_configuration_repository = agent_configuration_repository
         self._file_repository = file_repository
+        self._credentials_repository = credentials_repository
 
     def reset_agent_configuration(self):
         # NOTE: This method will be replaced by an event when we implement pub/sub in the island.
@@ -27,14 +34,9 @@ class RepositoryService:
         if custom_pbas.windows_filename:
             self._file_repository.delete_file(custom_pbas.windows_filename)
 
-    def unlock(self):
-        raise NotImplementedError
-
-    def reset_key(self):
-        raise NotImplementedError
-
     def clear_simulation_data(self):
         # NOTE: This method will be replaced by an event when we implement pub/sub in the island.
         #       Different plugins and components will be able to register for the event and clear
         #       any configuration data they've collected.
-        raise NotImplementedError
+        Database.reset_db(reset_config=False)
+        self._credentials_repository.remove_stolen_credentials()

monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py

@@ -5,7 +5,11 @@ import pytest
 from tests.monkey_island import InMemoryAgentConfigurationRepository
 
 from common.configuration import AgentConfiguration
-from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
+from monkey_island.cc.repository import (
+    IAgentConfigurationRepository,
+    ICredentialsRepository,
+    IFileRepository,
+)
 from monkey_island.cc.services import RepositoryService
 
 LINUX_FILENAME = "linux_pba_file.sh"
@@ -13,7 +17,7 @@ WINDOWS_FILENAME = "windows_pba_file.ps1"
 
 
 @pytest.fixture
-def agent_configuration(default_agent_configuration) -> AgentConfiguration:
+def agent_configuration(default_agent_configuration: AgentConfiguration) -> AgentConfiguration:
     custom_pbas = replace(
         default_agent_configuration.custom_pbas,
         linux_filename=LINUX_FILENAME,
@@ -23,7 +27,9 @@ def agent_configuration(default_agent_configuration) -> AgentConfiguration:
 
 
 @pytest.fixture
-def agent_configuration_repository(agent_configuration) -> IAgentConfigurationRepository:
+def agent_configuration_repository(
+    agent_configuration: AgentConfiguration,
+) -> IAgentConfigurationRepository:
     agent_configuration_repository = InMemoryAgentConfigurationRepository()
     agent_configuration_repository.store_configuration(agent_configuration)
 
@@ -31,15 +37,25 @@ def agent_configuration_repository(agent_configuration) -> IAgentConfigurationRe
 
 
 @pytest.fixture
-def mock_file_repository():
+def mock_file_repository() -> IFileRepository:
     return MagicMock(spec=IFileRepository)
 
 
-def test_reset_configuration__remove_pba_files(
-    agent_configuration_repository, mock_file_repository
-):
-    repository_service = RepositoryService(agent_configuration_repository, mock_file_repository)
+@pytest.fixture
+def mock_credentials_repository() -> ICredentialsRepository:
+    return MagicMock(spec=ICredentialsRepository)
 
+
+@pytest.fixture
+def repository_service(
+    agent_configuration_repository, mock_file_repository, mock_credentials_repository
+) -> RepositoryService:
+    return RepositoryService(
+        agent_configuration_repository, mock_file_repository, mock_credentials_repository
+    )
+
+
+def test_reset_configuration__remove_pba_files(repository_service, mock_file_repository):
     repository_service.reset_agent_configuration()
 
     assert mock_file_repository.delete_file.called_with(LINUX_FILENAME)
@@ -47,11 +63,20 @@ def test_reset_configuration__remove_pba_files(
 
 
 def test_reset_configuration__agent_configuration_changed(
-    agent_configuration_repository, agent_configuration, mock_file_repository
+    repository_service, agent_configuration_repository, agent_configuration
 ):
-    mock_file_repository = MagicMock(spec=IFileRepository)
-    repository_service = RepositoryService(agent_configuration_repository, mock_file_repository)
-
     repository_service.reset_agent_configuration()
 
     assert agent_configuration_repository.get_configuration() != agent_configuration
+
+
+@pytest.mark.usefixtures("uses_database")
+def test_clear_simulation_data(
+    repository_service: RepositoryService,
+    mock_credentials_repository: ICredentialsRepository,
+    monkeypatch,
+):
+    monkeypatch.setattr("monkey_island.cc.services.repository_service.Database", MagicMock())
+    repository_service.clear_simulation_data()
+
+    mock_credentials_repository.remove_stolen_credentials.assert_called_once()