Merge pull request #2332 from guardicore/2261-agent-signals-model

2261 agent signals model

monkey/monkey_island/cc/event_queue/i_island_event_queue.py

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

monkey/monkey_island/cc/models/__init__.py

@@ -15,3 +15,4 @@ from .communication_type import CommunicationType
 from .node import Node
 from common.types import AgentID
 from .agent import Agent
+from .agent_signals import AgentSignals

monkey/monkey_island/cc/models/agent_signals.py

@@ -0,0 +1,8 @@
+from datetime import datetime
+from typing import Optional
+
+from common.base_models import InfectionMonkeyBaseModel
+
+
+class AgentSignals(InfectionMonkeyBaseModel):
+    terminate: Optional[datetime]

monkey/monkey_island/cc/models/simulation.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+from datetime import datetime
 from enum import Enum
+from typing import Optional
 
 from common.base_models import InfectionMonkeyBaseModel
 
@@ -13,3 +15,4 @@ class IslandMode(Enum):
 
 class Simulation(InfectionMonkeyBaseModel):
     mode: IslandMode = IslandMode.UNSET
+    terminate_signal_time: Optional[datetime] = None

monkey/monkey_island/cc/services/__init__.py

@@ -1,3 +1,4 @@
 from .authentication_service import AuthenticationService
 
 from .aws import AWSService
+from .agent_signals_service import AgentSignalsService

monkey/monkey_island/cc/services/agent_signals_service.py

@@ -0,0 +1,26 @@
+from datetime import datetime
+
+from monkey_island.cc.models import Agent, AgentSignals
+from monkey_island.cc.repository import ISimulationRepository
+
+
+class AgentSignalsService:
+    def __init__(self, simulation_repository: ISimulationRepository):
+        self._simulation_repository = simulation_repository
+
+    def get_signals(self, agent: Agent) -> AgentSignals:
+        """
+        Gets the signals sent to a particular agent
+
+        :param agent: The agent whose signals need to be retrieved
+        :return: Signals sent to the relevant agent
+        """
+        return AgentSignals(timestamp=datetime.now())
+
+    def on_terminate_agents_signal(self, timestamp: datetime):
+        """
+        Updates the simulation repository with the terminate signal's timestamp
+
+        :param timestamp: Timestamp of the terminate signal
+        """
+        pass

vulture_allowlist.py

@@ -8,15 +8,14 @@ from common.agent_configuration.agent_sub_configurations import (
     ScanTargetConfiguration,
 )
 from common.credentials import Credentials, LMHash, NTHash
-from common.types import SocketAddress
 from infection_monkey.exploit.log4shell_utils.ldap_server import LDAPServerFactory
 from monkey_island.cc.event_queue import IslandEventTopic, PyPubSubIslandEventQueue
-from monkey_island.cc.models import Report
+from monkey_island.cc.models import Report, Simulation
+from monkey_island.cc.models import AgentSignals, Report
 from monkey_island.cc.models.networkmap import Arc, NetworkMap
 from monkey_island.cc.repository import MongoAgentRepository, MongoMachineRepository
 from monkey_island.cc.repository.attack.IMitigationsRepository import IMitigationsRepository
 from monkey_island.cc.repository.i_agent_event_repository import IAgentEventRepository
-from monkey_island.cc.repository.i_agent_log_repository import IAgentLogRepository
 from monkey_island.cc.repository.i_agent_repository import IAgentRepository
 from monkey_island.cc.repository.i_attack_repository import IAttackRepository
 from monkey_island.cc.repository.i_config_repository import IConfigRepository
@@ -27,6 +26,7 @@ from monkey_island.cc.repository.i_simulation_repository import ISimulationRepos
 from monkey_island.cc.repository.ICredentials import ICredentialsRepository
 from monkey_island.cc.repository.zero_trust.IEventRepository import IEventRepository
 from monkey_island.cc.repository.zero_trust.IFindingRepository import IFindingRepository
+from monkey_island.cc.services import AgentSignalsService
 
 fake_monkey_dir_path  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
 set_os_linux  # unused variable (monkey/tests/infection_monkey/post_breach/actions/test_users_custom_pba.py:37)
@@ -302,11 +302,6 @@ IAgentEventRepository.get_events_by_type
 IAgentEventRepository.get_events_by_tag
 IAgentEventRepository.get_events_by_source
 
-# TODO: Remove once #2274 is closed
-IAgentLogRepository
-IAgentLogRepository.upsert_agent_log
-IAgentLogRepository.get_agent_log
-
 
 # pydantic base models
 underscore_attrs_are_private
@@ -326,5 +321,13 @@ EXPLOITED
 CC
 CC_TUNNEL
 
-# TODO: Remove after #2323
-SocketAddress
+IslandEventTopic.AGENT_CONNECTED
+IslandEventTopic.CLEAR_SIMULATION_DATA
+IslandEventTopic.RESET_AGENT_CONFIGURATION
+# TODO: Remove after #2261 is closed
+IslandEventTopic.TERMINATE_AGENTS
+
+Simulation.terminate_signal_time
+
+AgentSignalsService.get_signals
+AgentSignalsService.on_terminate_agents_signal