Island: Implement RepositoryService.reset_agent_configuration()

This commit is contained in:
Mike Salvatore 2022-07-13 08:54:53 -04:00
parent a4b8c94480
commit 726f2bbaa2
2 changed files with 83 additions and 1 deletions

View File

@ -1,6 +1,31 @@
from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
class RepositoryService:
def __init__(
self,
agent_configuration_repository: IAgentConfigurationRepository,
file_repository: IFileRepository,
):
self._agent_configuration_repository = agent_configuration_repository
self._file_repository = file_repository
def reset_agent_configuration(self):
pass
# 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 reset
# their configurations.
self._remove_pba_files()
self._agent_configuration_repository.reset_to_default()
def _remove_pba_files(self):
agent_configuration = self._agent_configuration_repository.get_configuration()
custom_pbas = agent_configuration.custom_pbas
if custom_pbas.linux_filename:
self._file_repository.delete_file(custom_pbas.linux_filename)
if custom_pbas.windows_filename:
self._file_repository.delete_file(custom_pbas.windows_filename)
def unlock(self):
pass

View File

@ -0,0 +1,57 @@
from dataclasses import replace
from unittest.mock import MagicMock
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.services import RepositoryService
LINUX_FILENAME = "linux_pba_file.sh"
WINDOWS_FILENAME = "windows_pba_file.ps1"
@pytest.fixture
def agent_configuration(default_agent_configuration) -> AgentConfiguration:
custom_pbas = replace(
default_agent_configuration.custom_pbas,
linux_filename=LINUX_FILENAME,
windows_filename=WINDOWS_FILENAME,
)
return replace(default_agent_configuration, custom_pbas=custom_pbas)
@pytest.fixture
def agent_configuration_repository(agent_configuration) -> IAgentConfigurationRepository:
agent_configuration_repository = InMemoryAgentConfigurationRepository()
agent_configuration_repository.store_configuration(agent_configuration)
return agent_configuration_repository
@pytest.fixture
def mock_file_repository():
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)
repository_service.reset_agent_configuration()
assert mock_file_repository.delete_file.called_with(LINUX_FILENAME)
assert mock_file_repository.delete_file.called_with(LINUX_FILENAME)
def test_reset_configuration__agent_configuration_changed(
agent_configuration_repository, agent_configuration, mock_file_repository
):
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