p15670423
/
monkey
1
0
Fork 5
Code Issues Pull Requests 3 Projects Releases Wiki Activity

Merge pull request #2327 from guardicore/2266-replace-IslandModeService 

Replace IslandModeService with event
This commit is contained in:
Mike Salvatore 2022-09-21 09:25:29 -04:00 committed by GitHub
parent 18ceb6c279 a8fb05f2b8
commit a76273fa0d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 102 additions and 144 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
monkey/monkey_island/cc/event_queue/i_island_event_queue.py
View File

@ -9,6 +9,7 @@ class IslandEventTopic(Enum):
AGENT_CONNECTED = auto()
CLEAR_SIMULATION_DATA = auto()
RESET_AGENT_CONFIGURATION = auto()
SET_ISLAND_MODE = auto()
class IIslandEventQueue(ABC):

1
monkey/monkey_island/cc/island_event_handlers/__init__.py
View File

@ -1,2 +1,3 @@
from .reset_agent_configuration import reset_agent_configuration
from .reset_machine_repository import reset_machine_repository
from .set_agent_configuration_per_island_mode import set_agent_configuration_per_island_mode

29
monkey/monkey_island/cc/island_event_handlers/set_agent_configuration_per_island_mode.py Normal file
View File

@ -0,0 +1,29 @@
from common.agent_configuration import AgentConfiguration
from monkey_island.cc.models import IslandMode
from monkey_island.cc.repository import IAgentConfigurationRepository
class set_agent_configuration_per_island_mode:
"""
Callable class that sets the default Agent configuration as per the Island's mode
"""
def __init__(
self,
agent_configuration_repository: IAgentConfigurationRepository,
default_agent_configuration: AgentConfiguration,
default_ransomware_agent_configuration: AgentConfiguration,
):
self._agent_configuration_repository = agent_configuration_repository
self._default_agent_configuration = default_agent_configuration
self._default_ransomware_agent_configuration = default_ransomware_agent_configuration
def __call__(self, mode: IslandMode):
if mode == IslandMode.RANSOMWARE:
self._agent_configuration_repository.store_configuration(
self._default_ransomware_agent_configuration
)
else:
self._agent_configuration_repository.store_configuration(
self._default_agent_configuration
)

18
monkey/monkey_island/cc/resources/island_mode.py
View File

@ -4,10 +4,11 @@ from http import HTTPStatus
from flask import request
from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
from monkey_island.cc.models import IslandMode as IslandModeEnum
from monkey_island.cc.repository import ISimulationRepository
from monkey_island.cc.resources.AbstractResource import AbstractResource
from monkey_island.cc.resources.request_authentication import jwt_required
from monkey_island.cc.services import IslandModeService
logger = logging.getLogger(__name__)
@ -15,16 +16,19 @@ logger = logging.getLogger(__name__)
class IslandMode(AbstractResource):
urls = ["/api/island/mode"]
def __init__(self, island_mode_service: IslandModeService):
self._island_mode_service = island_mode_service
def __init__(
self,
island_event_queue: IIslandEventQueue,
simulation_repository: ISimulationRepository,
):
self._island_event_queue = island_event_queue
self._simulation_repository = simulation_repository
@jwt_required
def put(self):
try:
mode = IslandModeEnum(request.json)
self._island_mode_service.set_mode(mode)
self._island_event_queue.publish(topic=IslandEventTopic.SET_ISLAND_MODE, mode=mode)
return {}, HTTPStatus.NO_CONTENT
except (AttributeError, json.decoder.JSONDecodeError):
return {}, HTTPStatus.BAD_REQUEST
@ -33,5 +37,5 @@ class IslandMode(AbstractResource):
@jwt_required
def get(self):
island_mode = self._island_mode_service.get_mode()
island_mode = self._simulation_repository.get_mode()
return island_mode.value, HTTPStatus.OK

1
monkey/monkey_island/cc/services/__init__.py
View File

@ -1,4 +1,3 @@
from .authentication_service import AuthenticationService
from .aws import AWSService
from .island_mode_service import IslandModeService

3
monkey/monkey_island/cc/services/initialize.py
View File

@ -47,7 +47,7 @@ from monkey_island.cc.repository import (
)
from monkey_island.cc.server_utils.consts import MONKEY_ISLAND_ABS_PATH
from monkey_island.cc.server_utils.encryption import ILockableEncryptor, RepositoryEncryptor
from monkey_island.cc.services import AWSService, IslandModeService
from monkey_island.cc.services import AWSService
from monkey_island.cc.services.attack.technique_reports.T1003 import T1003, T1003GetReportData
from monkey_island.cc.services.run_local_monkey import LocalMonkeyRunService
from monkey_island.cc.setup.mongo.mongo_setup import MONGO_URL
@ -175,7 +175,6 @@ def _log_agent_binary_hashes(agent_binary_repository: IAgentBinaryRepository):
def _register_services(container: DIContainer):
container.register_instance(AWSService, container.resolve(AWSService))
container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService))
container.register_instance(IslandModeService, container.resolve(IslandModeService))
container.register_instance(AuthenticationService, container.resolve(AuthenticationService))

46
monkey/monkey_island/cc/services/island_mode_service.py
View File

@ -1,46 +0,0 @@
from common.agent_configuration import AgentConfiguration
from monkey_island.cc.models import IslandMode
from monkey_island.cc.repository import IAgentConfigurationRepository, ISimulationRepository
class IslandModeService:
def __init__(
self,
_agent_configuration_repository: IAgentConfigurationRepository,
simulation_repository: ISimulationRepository,
default_agent_configuration: AgentConfiguration,
default_ransomware_agent_configuration: AgentConfiguration,
):
self._agent_configuration_repository = _agent_configuration_repository
self._simulation_repository = simulation_repository
self._default_agent_configuration = default_agent_configuration
self._default_ransomware_agent_configuration = default_ransomware_agent_configuration
def get_mode(self):
"""
Get's the island's current mode
:return The island's current mode
:raises RetrievalError: If the mode could not be retrieved
"""
return self._simulation_repository.get_mode()
def set_mode(self, mode: IslandMode):
"""
Set the island's mode
:param mode: The island's new mode
:raises StorageError: If the mode could not be saved
"""
self._simulation_repository.set_mode(mode)
self._set_configuration(mode)
def _set_configuration(self, mode: IslandMode):
if mode == IslandMode.RANSOMWARE:
self._agent_configuration_repository.store_configuration(
self._default_ransomware_agent_configuration
)
else:
self._agent_configuration_repository.store_configuration(
self._default_agent_configuration
)

34
monkey/monkey_island/cc/setup/island_event_handlers.py
View File

@ -5,12 +5,14 @@ from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic
from monkey_island.cc.island_event_handlers import (
reset_agent_configuration,
reset_machine_repository,
set_agent_configuration_per_island_mode,
)
from monkey_island.cc.repository import (
IAgentEventRepository,
IAgentRepository,
ICredentialsRepository,
INodeRepository,
ISimulationRepository,
)
from monkey_island.cc.services.database import Database
@ -20,30 +22,29 @@ def setup_island_event_handlers(container: DIContainer):
_subscribe_reset_agent_configuration_events(island_event_queue, container)
_subscribe_clear_simulation_data_events(island_event_queue, container)
_subscribe_set_island_mode_events(island_event_queue, container)
def _subscribe_reset_agent_configuration_events(
island_event_queue: IIslandEventQueue, container: DIContainer
):
island_event_queue.subscribe(
IslandEventTopic.RESET_AGENT_CONFIGURATION, container.resolve(reset_agent_configuration)
)
topic = IslandEventTopic.RESET_AGENT_CONFIGURATION
island_event_queue.subscribe(topic, container.resolve(reset_agent_configuration))
def _subscribe_clear_simulation_data_events(
island_event_queue: IIslandEventQueue, container: DIContainer
):
topic = IslandEventTopic.CLEAR_SIMULATION_DATA
legacy_database_reset = partial(Database.reset_db, reset_config=False)
island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, legacy_database_reset)
island_event_queue.subscribe(topic, legacy_database_reset)
credentials_repository = container.resolve(ICredentialsRepository)
island_event_queue.subscribe(
IslandEventTopic.CLEAR_SIMULATION_DATA, credentials_repository.remove_stolen_credentials
)
island_event_queue.subscribe(topic, credentials_repository.remove_stolen_credentials)
island_event_queue.subscribe(
IslandEventTopic.CLEAR_SIMULATION_DATA, container.resolve(reset_machine_repository)
)
island_event_queue.subscribe(topic, container.resolve(reset_machine_repository))
for i_repository in [
INodeRepository,
@ -51,4 +52,15 @@ def _subscribe_clear_simulation_data_events(
IAgentRepository,
]:
repository = container.resolve(i_repository)
island_event_queue.subscribe(IslandEventTopic.CLEAR_SIMULATION_DATA, repository.reset)
island_event_queue.subscribe(topic, repository.reset)
def _subscribe_set_island_mode_events(
island_event_queue: IIslandEventQueue, container: DIContainer
):
topic = IslandEventTopic.SET_ISLAND_MODE
island_event_queue.subscribe(topic, container.resolve(set_agent_configuration_per_island_mode))
simulation_repository = container.resolve(ISimulationRepository)
island_event_queue.subscribe(topic, simulation_repository.set_mode)

68
monkey/tests/unit_tests/monkey_island/cc/resources/test_island_mode.py
View File

@ -5,30 +5,42 @@ import pytest
from tests.common import StubDIContainer
from tests.monkey_island import InMemorySimulationRepository
from monkey_island.cc.event_queue import IIslandEventQueue
from monkey_island.cc.models import IslandMode
from monkey_island.cc.repository import RetrievalError
from monkey_island.cc.repository import ISimulationRepository
from monkey_island.cc.resources.island_mode import IslandMode as IslandModeResource
from monkey_island.cc.services import IslandModeService
class MockIslandModeService(IslandModeService):
def __init__(self):
self._simulation_repository = InMemorySimulationRepository()
def get_mode(self) -> IslandMode:
return self._simulation_repository.get_mode()
def set_mode(self, mode: IslandMode):
self._simulation_repository.set_mode(mode)
@pytest.fixture
def flask_client(build_flask_client):
container = StubDIContainer()
container.register_instance(IslandModeService, MockIslandModeService())
def flask_client_builder(build_flask_client):
def inner(side_effect=None):
container = StubDIContainer()
with build_flask_client(container) as flask_client:
yield flask_client
in_memory_simulation_repository = InMemorySimulationRepository()
container.register_instance(ISimulationRepository, in_memory_simulation_repository)
mock_island_event_queue = MagicMock(spec=IIslandEventQueue)
mock_island_event_queue.publish.side_effect = (
side_effect
if side_effect
else lambda topic, mode: in_memory_simulation_repository.set_mode(mode)
)
container.register_instance(IIslandEventQueue, mock_island_event_queue)
with build_flask_client(container) as flask_client:
return flask_client
return inner
@pytest.fixture
def flask_client(flask_client_builder):
return flask_client_builder()
@pytest.fixture
def flask_client__internal_server_error(flask_client_builder):
return flask_client_builder(Exception)
@pytest.mark.parametrize(
@ -53,20 +65,12 @@ def test_island_mode_post__invalid_mode(flask_client):
assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
def test_island_mode_post__internal_server_error(build_flask_client):
mock_island_mode_service = MagicMock(spec=IslandModeService)
mock_island_mode_service.set_mode = MagicMock(side_effect=RetrievalError)
container = StubDIContainer()
container.register_instance(IslandModeService, mock_island_mode_service)
with build_flask_client(container) as flask_client:
resp = flask_client.put(
IslandModeResource.urls[0],
json=IslandMode.RANSOMWARE.value,
follow_redirects=True,
)
def test_island_mode_post__internal_server_error(flask_client__internal_server_error):
resp = flask_client__internal_server_error.put(
IslandModeResource.urls[0],
json=IslandMode.RANSOMWARE.value,
follow_redirects=True,
)
assert resp.status_code == HTTPStatus.INTERNAL_SERVER_ERROR

45
monkey/tests/unit_tests/monkey_island/cc/services/test_island_mode_service.py
View File

@ -1,45 +0,0 @@
import pytest
from tests.monkey_island import InMemoryAgentConfigurationRepository, InMemorySimulationRepository
from common.agent_configuration import (
DEFAULT_AGENT_CONFIGURATION,
DEFAULT_RANSOMWARE_AGENT_CONFIGURATION,
)
from monkey_island.cc.models import IslandMode
from monkey_island.cc.services import IslandModeService
@pytest.fixture
def agent_configuration_repository():
return InMemoryAgentConfigurationRepository()
@pytest.fixture
def island_mode_service(agent_configuration_repository):
return IslandModeService(
agent_configuration_repository,
InMemorySimulationRepository(),
DEFAULT_AGENT_CONFIGURATION,
DEFAULT_RANSOMWARE_AGENT_CONFIGURATION,
)
@pytest.mark.parametrize("mode", list(IslandMode))
def test_set_mode(island_mode_service, mode):
island_mode_service.set_mode(mode)
assert island_mode_service.get_mode() == mode
@pytest.mark.parametrize(
"mode, expected_config",
[
(IslandMode.UNSET, DEFAULT_AGENT_CONFIGURATION),
(IslandMode.ADVANCED, DEFAULT_AGENT_CONFIGURATION),
(IslandMode.RANSOMWARE, DEFAULT_RANSOMWARE_AGENT_CONFIGURATION),
],
)
def test_set_mode_sets_config(
island_mode_service, agent_configuration_repository, mode, expected_config
):
island_mode_service.set_mode(mode)
assert agent_configuration_repository.get_configuration() == expected_config