From f12e8398787b8cf2f0078a7f688f876c08ac032e Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 17:34:57 +0200 Subject: [PATCH 1/5] Island: Register AgentSignalsSerivce in DI Container --- monkey/monkey_island/cc/services/initialize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index 830e010d7..c0dc600c3 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -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 +from monkey_island.cc.services import AgentSignalsService, 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 @@ -176,6 +176,7 @@ def _register_services(container: DIContainer): container.register_instance(AWSService, container.resolve(AWSService)) container.register_instance(LocalMonkeyRunService, container.resolve(LocalMonkeyRunService)) container.register_instance(AuthenticationService, container.resolve(AuthenticationService)) + container.register_instance(AgentSignalsService, container.resolve(AgentSignalsService)) def _dirty_hacks(container: DIContainer): From 5eeee2a60d7bf2873f02bb2b47d38e3f559b58a9 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 17:36:42 +0200 Subject: [PATCH 2/5] Island: Subscribe AgentSignalsService.on_terminate_agents_signal to TERMINATE_AGENTS events --- .../monkey_island/cc/setup/island_event_handlers.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 6b4cb8c53..31e7908f6 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -15,6 +15,7 @@ from monkey_island.cc.repository import ( INodeRepository, ISimulationRepository, ) +from monkey_island.cc.services import AgentSignalsService from monkey_island.cc.services.database import Database @@ -25,6 +26,7 @@ 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) + _subscribe_on_terminate_agents_signal(island_event_queue, container) def _subscribe_agent_registration_events( @@ -74,3 +76,13 @@ def _subscribe_set_island_mode_events( simulation_repository = container.resolve(ISimulationRepository) island_event_queue.subscribe(topic, simulation_repository.set_mode) + + +def _subscribe_on_terminate_agents_signal( + island_event_queue: IIslandEventQueue, container: DIContainer +): + topic = IslandEventTopic.TERMINATE_AGENTS + + agent_signals_service = container.resolve(AgentSignalsService) + + island_event_queue.subscribe(topic, agent_signals_service.on_terminate_agents_signal) From 9d3be7e1d3936fff758d94fa8e05e3bfa15112e2 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 22 Sep 2022 17:37:14 +0200 Subject: [PATCH 3/5] Island: Implement AgentSignalsService.on_terminate_agents_signal --- .../monkey_island/cc/services/agent_signals_service.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index 23a15dba9..473a10066 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -1,10 +1,13 @@ +import logging from datetime import datetime from typing import Optional from common.types import AgentID -from monkey_island.cc.models import AgentSignals +from monkey_island.cc.models import AgentSignals, Simulation from monkey_island.cc.repository import IAgentRepository, ISimulationRepository +logger = logging.getLogger(__name__) + class AgentSignalsService: def __init__( @@ -45,4 +48,7 @@ class AgentSignalsService: :param timestamp: Timestamp of the terminate signal """ - pass + simulation = self._simulation_repository.get_simulation() + updated_simulation = Simulation(mode=simulation.mode, terminate_signal_time=timestamp) + + self._simulation_repository.save_simulation(updated_simulation) From 275efb2ab1531c34ed7ed98a1dde04ea7eda5f2f Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Thu, 22 Sep 2022 19:23:36 +0000 Subject: [PATCH 4/5] UT: Test on_terminate_agents_signal --- monkey/monkey_island/cc/services/__init__.py | 2 +- .../cc/services/test_agent_signals_service.py | 27 ++++++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/__init__.py b/monkey/monkey_island/cc/services/__init__.py index c73aff356..8fdb5fc77 100644 --- a/monkey/monkey_island/cc/services/__init__.py +++ b/monkey/monkey_island/cc/services/__init__.py @@ -1,4 +1,4 @@ +from .agent_signals_service import AgentSignalsService from .authentication_service import AuthenticationService from .aws import AWSService -from .agent_signals_service import AgentSignalsService diff --git a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py index 16aab943f..1beb53940 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py +++ b/monkey/tests/unit_tests/monkey_island/cc/services/test_agent_signals_service.py @@ -4,7 +4,7 @@ from uuid import UUID import pytest from common.types import AgentID -from monkey_island.cc.models import Agent, Simulation +from monkey_island.cc.models import Agent, IslandMode, Simulation from monkey_island.cc.repository import IAgentRepository, ISimulationRepository, UnknownRecordError from monkey_island.cc.services import AgentSignalsService @@ -117,3 +117,28 @@ def test_progenitor_started_before_terminate( signals = agent_signals_service.get_signals(agent.id) assert signals.terminate.timestamp() == TERMINATE_TIMESTAMP + + +def test_on_terminate_agents_signal__stores_timestamp( + agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository +): + timestamp = 100 + mock_simulation_repository.get_simulation = MagicMock(return_value=Simulation()) + agent_signals_service.on_terminate_agents_signal(timestamp) + + expected_value = Simulation(terminate_signal_time=timestamp) + assert mock_simulation_repository.save_simulation.called_once_with(expected_value) + + +def test_on_terminate_agents_signal__updates_timestamp( + agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository +): + timestamp = 100 + mock_simulation_repository.get_simulation = MagicMock( + return_value=Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=50) + ) + + agent_signals_service.on_terminate_agents_signal(timestamp) + + expected_value = Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=timestamp) + assert mock_simulation_repository.save_simulation.called_once_with(expected_value) From f9306cf8f157802998986a0e9ac19fd6d6396494 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Fri, 23 Sep 2022 09:41:31 +0200 Subject: [PATCH 5/5] Island: Keep naming consistency in island_event_handlers --- monkey/monkey_island/cc/setup/island_event_handlers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 31e7908f6..96d2e1941 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -26,7 +26,7 @@ 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) - _subscribe_on_terminate_agents_signal(island_event_queue, container) + _subscribe_terminate_agents_events(island_event_queue, container) def _subscribe_agent_registration_events( @@ -78,7 +78,7 @@ def _subscribe_set_island_mode_events( island_event_queue.subscribe(topic, simulation_repository.set_mode) -def _subscribe_on_terminate_agents_signal( +def _subscribe_terminate_agents_events( island_event_queue: IIslandEventQueue, container: DIContainer ): topic = IslandEventTopic.TERMINATE_AGENTS