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/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) 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): diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index 6b4cb8c53..96d2e1941 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_terminate_agents_events(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_terminate_agents_events( + 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) 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)