forked from p15670423/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:
|
||||
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
|
||||
|
|
|
@ -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