From d2db1d167373d66de053bf7c24dca67a83e21631 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 6 Mar 2023 21:24:12 +0000 Subject: [PATCH 1/4] Island: Set terminate signal for duplicate agents If more than one agent is running on a given machine, the first agent shall be allowed to run, and the rest shall be given a terminate signal. --- .../cc/services/agent_signals_service.py | 20 ++++++++- .../cc/services/test_agent_signals_service.py | 43 ++++++++++++++----- 2 files changed, 52 insertions(+), 11 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index b47ddcb62e6..b583b93be16 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -4,7 +4,7 @@ from common.agent_signals import AgentSignals from common.types import AgentID -from monkey_island.cc.models import Simulation, TerminateAllAgents +from monkey_island.cc.models import Agent, MachineID, Simulation, TerminateAllAgents from monkey_island.cc.repositories import IAgentRepository, ISimulationRepository logger = logging.getLogger(__name__) @@ -40,9 +40,27 @@ def get_signals(self, agent_id: AgentID) -> AgentSignals: if agent.stop_time is not None: return AgentSignals(terminate=agent.stop_time) + if not self._agent_is_first_to_register(agent): + return AgentSignals(terminate=agent.registration_time) + terminate_timestamp = self._get_terminate_signal_timestamp(agent_id) return AgentSignals(terminate=terminate_timestamp) + def _agents_running_on_machine(self, machine_id: MachineID): + return [ + a + for a in self._agent_repository.get_agents() + if a.machine_id == machine_id and a.stop_time is None + ] + + def _agent_is_first_to_register(self, agent: Agent) -> bool: + agents_on_same_machine = self._agents_running_on_machine(agent.machine_id) + print(agents_on_same_machine) + first_to_register = min( + agents_on_same_machine, key=lambda a: a.registration_time, default=agent + ) + return agent is first_to_register + def _get_terminate_signal_timestamp(self, agent_id: AgentID) -> Optional[datetime]: simulation = self._simulation_repository.get_simulation() terminate_all_signal_time = simulation.terminate_signal_time 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 7f634ffae7f..8fc2027affe 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 @@ -29,10 +29,19 @@ AGENT_3 = Agent( id=UUID("0fc9afcb-1902-436b-bd5c-1ad194252484"), machine_id=3, + registration_time=301, start_time=300, parent_id=AGENT_2.id, ) +AGENT_4 = Agent( + id=UUID("0fc9afcb-1902-436b-bd5c-1ad194252485"), + machine_id=3, + registration_time=302, + start_time=299, + parent_id=AGENT_2.id, +) + AGENTS = [AGENT_1, AGENT_2, AGENT_3] STOPPED_AGENT = Agent( @@ -43,12 +52,14 @@ parent_id=AGENT_3.id, ) -ALL_AGENTS = [*AGENTS, STOPPED_AGENT] +ALL_AGENTS = [*AGENTS, AGENT_4, STOPPED_AGENT] @pytest.fixture def mock_simulation_repository() -> IAgentRepository: - return MagicMock(spec=ISimulationRepository) + repository = MagicMock(spec=ISimulationRepository) + repository.get_simulation = MagicMock(return_value=Simulation(terminate_signal_time=None)) + return repository @pytest.fixture(scope="session") @@ -63,6 +74,7 @@ def get_agent_by_id(agent_id: AgentID) -> Agent: agent_repository = MagicMock(spec=IAgentRepository) agent_repository.get_progenitor = MagicMock(return_value=AGENT_1) agent_repository.get_agent_by_id = MagicMock(side_effect=get_agent_by_id) + agent_repository.get_agents = MagicMock(return_value=ALL_AGENTS) return agent_repository @@ -77,9 +89,6 @@ def test_stopped_agent( mock_simulation_repository: ISimulationRepository, ): agent = STOPPED_AGENT - mock_simulation_repository.get_simulation = MagicMock( - return_value=Simulation(terminate_signal_time=None) - ) signals = agent_signals_service.get_signals(agent.id) assert signals.terminate == agent.stop_time @@ -91,10 +100,6 @@ def test_terminate_is_none( agent_signals_service: AgentSignalsService, mock_simulation_repository: ISimulationRepository, ): - mock_simulation_repository.get_simulation = MagicMock( - return_value=Simulation(terminate_signal_time=None) - ) - signals = agent_signals_service.get_signals(agent.id) assert signals.terminate is None @@ -153,7 +158,6 @@ def test_on_terminate_agents_signal__stores_timestamp( timestamp = 100 terminate_all_agents = TerminateAllAgents(timestamp=timestamp) - mock_simulation_repository.get_simulation = MagicMock(return_value=Simulation()) agent_signals_service.on_terminate_agents_signal(terminate_all_agents) expected_value = Simulation(terminate_signal_time=timestamp) @@ -174,3 +178,22 @@ def test_on_terminate_agents_signal__updates_timestamp( expected_value = Simulation(mode=IslandMode.RANSOMWARE, terminate_signal_time=timestamp) assert mock_simulation_repository.save_simulation.called_once_with(expected_value) + + +def test_terminate_signal__not_set_if_agent_registered_before_another(agent_signals_service): + signals = agent_signals_service.get_signals(AGENT_3.id) + + assert signals.terminate is None + + +def test_terminate_signal__set_if_agent_registered_after_another(agent_signals_service): + signals = agent_signals_service.get_signals(AGENT_4.id) + + assert signals.terminate is not None + + +def test_terminate_signal__not_set_if_agent_registered_after_stopped_agent(agent_signals_service): + AGENT_3.stop_time = 400 + signals = agent_signals_service.get_signals(AGENT_4.id) + + assert signals.terminate is None From 61e4da3289711814362c1afb7202417895a78854 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Tue, 7 Mar 2023 11:17:25 +0100 Subject: [PATCH 2/4] Island: Change running_agents_on_machine logic in AgentSignalsService --- .../cc/services/agent_signals_service.py | 13 +++++-------- .../cc/services/test_agent_signals_service.py | 8 +++++--- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/monkey/monkey_island/cc/services/agent_signals_service.py b/monkey/monkey_island/cc/services/agent_signals_service.py index b583b93be16..e58b21cb4eb 100644 --- a/monkey/monkey_island/cc/services/agent_signals_service.py +++ b/monkey/monkey_island/cc/services/agent_signals_service.py @@ -46,21 +46,18 @@ def get_signals(self, agent_id: AgentID) -> AgentSignals: terminate_timestamp = self._get_terminate_signal_timestamp(agent_id) return AgentSignals(terminate=terminate_timestamp) - def _agents_running_on_machine(self, machine_id: MachineID): - return [ - a - for a in self._agent_repository.get_agents() - if a.machine_id == machine_id and a.stop_time is None - ] - def _agent_is_first_to_register(self, agent: Agent) -> bool: agents_on_same_machine = self._agents_running_on_machine(agent.machine_id) - print(agents_on_same_machine) first_to_register = min( agents_on_same_machine, key=lambda a: a.registration_time, default=agent ) return agent is first_to_register + def _agents_running_on_machine(self, machine_id: MachineID): + return [ + a for a in self._agent_repository.get_running_agents() if a.machine_id == machine_id + ] + def _get_terminate_signal_timestamp(self, agent_id: AgentID) -> Optional[datetime]: simulation = self._simulation_repository.get_simulation() terminate_all_signal_time = simulation.terminate_signal_time 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 8fc2027affe..76746d582e3 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 @@ -74,7 +74,7 @@ def get_agent_by_id(agent_id: AgentID) -> Agent: agent_repository = MagicMock(spec=IAgentRepository) agent_repository.get_progenitor = MagicMock(return_value=AGENT_1) agent_repository.get_agent_by_id = MagicMock(side_effect=get_agent_by_id) - agent_repository.get_agents = MagicMock(return_value=ALL_AGENTS) + agent_repository.get_running_agents = MagicMock(return_value=AGENTS) return agent_repository @@ -192,8 +192,10 @@ def test_terminate_signal__set_if_agent_registered_after_another(agent_signals_s assert signals.terminate is not None -def test_terminate_signal__not_set_if_agent_registered_after_stopped_agent(agent_signals_service): - AGENT_3.stop_time = 400 +def test_terminate_signal__not_set_if_agent_registered_after_stopped_agent( + agent_signals_service: AgentSignalsService, mock_agent_repository: IAgentRepository +): + mock_agent_repository.get_running_agents = MagicMock(return_value=[AGENT_1, AGENT_2]) signals = agent_signals_service.get_signals(AGENT_4.id) assert signals.terminate is None From 651b4e8c19bf1bd43d4a578b5481190098fcd65d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 7 Mar 2023 08:18:27 -0500 Subject: [PATCH 3/4] UT: Improve implementation of mock_agent_repository.get_running_agents --- .../monkey_island/cc/services/test_agent_signals_service.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 76746d582e3..b0498f71489 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 @@ -74,7 +74,9 @@ def get_agent_by_id(agent_id: AgentID) -> Agent: agent_repository = MagicMock(spec=IAgentRepository) agent_repository.get_progenitor = MagicMock(return_value=AGENT_1) agent_repository.get_agent_by_id = MagicMock(side_effect=get_agent_by_id) - agent_repository.get_running_agents = MagicMock(return_value=AGENTS) + agent_repository.get_running_agents = MagicMock( + return_value=[a for a in ALL_AGENTS if a.stop_time is None] + ) return agent_repository From 3c66438f2f7fa6f07871609a885963daa73e55db Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 7 Mar 2023 08:20:35 -0500 Subject: [PATCH 4/4] UT: Rename AGENT_4 -> DUPLICATE_MACHINE_AGENT --- .../cc/services/test_agent_signals_service.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 b0498f71489..bf0edd3782d 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 @@ -34,7 +34,7 @@ parent_id=AGENT_2.id, ) -AGENT_4 = Agent( +DUPLICATE_MACHINE_AGENT = Agent( id=UUID("0fc9afcb-1902-436b-bd5c-1ad194252485"), machine_id=3, registration_time=302, @@ -52,7 +52,7 @@ parent_id=AGENT_3.id, ) -ALL_AGENTS = [*AGENTS, AGENT_4, STOPPED_AGENT] +ALL_AGENTS = [*AGENTS, DUPLICATE_MACHINE_AGENT, STOPPED_AGENT] @pytest.fixture @@ -189,7 +189,7 @@ def test_terminate_signal__not_set_if_agent_registered_before_another(agent_sign def test_terminate_signal__set_if_agent_registered_after_another(agent_signals_service): - signals = agent_signals_service.get_signals(AGENT_4.id) + signals = agent_signals_service.get_signals(DUPLICATE_MACHINE_AGENT.id) assert signals.terminate is not None @@ -198,6 +198,6 @@ def test_terminate_signal__not_set_if_agent_registered_after_stopped_agent( agent_signals_service: AgentSignalsService, mock_agent_repository: IAgentRepository ): mock_agent_repository.get_running_agents = MagicMock(return_value=[AGENT_1, AGENT_2]) - signals = agent_signals_service.get_signals(AGENT_4.id) + signals = agent_signals_service.get_signals(DUPLICATE_MACHINE_AGENT.id) assert signals.terminate is None