forked from p34709852/monkey
Island: Implement RepositoryService.reset_agent_configuration()
This commit is contained in:
parent
a4b8c94480
commit
726f2bbaa2
|
@ -1,6 +1,31 @@
|
||||||
|
from monkey_island.cc.repository import IAgentConfigurationRepository, IFileRepository
|
||||||
|
|
||||||
|
|
||||||
class RepositoryService:
|
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):
|
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):
|
def unlock(self):
|
||||||
pass
|
pass
|
||||||
|
|
|
@ -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
|
Loading…
Reference in New Issue