From 37885b786f460e0b806462e2b79989be83f033fc Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 11:09:46 -0400 Subject: [PATCH 1/4] UT: Add missing type hints to test_repository_service --- .../monkey_island/cc/services/test_repository_service.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py index adc07eb2e..17dd6ffb5 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py @@ -13,7 +13,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 +23,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,7 +33,7 @@ def agent_configuration_repository(agent_configuration) -> IAgentConfigurationRe @pytest.fixture -def mock_file_repository(): +def mock_file_repository() -> IFileRepository: return MagicMock(spec=IFileRepository) From fcff724eaaa0e401862fadd825ccd7c3cd530741 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 11:25:00 -0400 Subject: [PATCH 2/4] Island: Implement RepositoryService.clear_simulation_data() --- .../cc/services/repository_service.py | 12 +++++- .../cc/services/test_repository_service.py | 41 +++++++++++++++---- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/monkey/monkey_island/cc/services/repository_service.py b/monkey/monkey_island/cc/services/repository_service.py index af0d13fa1..9910469b5 100644 --- a/monkey/monkey_island/cc/services/repository_service.py +++ b/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. @@ -37,4 +44,5 @@ class RepositoryService: # 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() diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py index 17dd6ffb5..a5d443456 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_repository_service.py +++ b/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" @@ -37,11 +41,21 @@ 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) @@ -49,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() From 6ca09d5c9466639c7967abfbe404be48bf33a2e0 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 11:25:18 -0400 Subject: [PATCH 3/4] Island: Use RepositoryService in ClearSimulationData resource --- .../cc/resources/clear_simulation_data.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/monkey/monkey_island/cc/resources/clear_simulation_data.py b/monkey/monkey_island/cc/resources/clear_simulation_data.py index fa85cac70..342fadd5f 100644 --- a/monkey/monkey_island/cc/resources/clear_simulation_data.py +++ b/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) From 6e60722adc5ace4f124892972797b8f102b7b0a6 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 13 Jul 2022 19:35:19 -0400 Subject: [PATCH 4/4] Island: Remove unlock() and reset_key() from RepositoryService These two methods violate SRP and ISP. --- monkey/monkey_island/cc/services/repository_service.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/monkey/monkey_island/cc/services/repository_service.py b/monkey/monkey_island/cc/services/repository_service.py index 9910469b5..8cc310b62 100644 --- a/monkey/monkey_island/cc/services/repository_service.py +++ b/monkey/monkey_island/cc/services/repository_service.py @@ -34,12 +34,6 @@ 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