Merge pull request #2337 from guardicore/2261-on-terminate-agents-signal

2261 on terminate agents signal

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

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)

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):

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)

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)