From d514ac283c8ae6eb74489ea27ea158442c229823 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 20 Sep 2022 15:01:00 -0400 Subject: [PATCH 01/12] Island: Publish AgentRegistrationData from Agents resource --- monkey/monkey_island/cc/resources/agents.py | 7 +++++++ .../monkey_island/cc/resources/test_agents.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/monkey/monkey_island/cc/resources/agents.py b/monkey/monkey_island/cc/resources/agents.py index b727bd9ae..4fb9837ff 100644 --- a/monkey/monkey_island/cc/resources/agents.py +++ b/monkey/monkey_island/cc/resources/agents.py @@ -5,6 +5,7 @@ from http import HTTPStatus from flask import make_response, request from common import AgentRegistrationData +from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic from monkey_island.cc.resources.AbstractResource import AbstractResource logger = logging.getLogger(__name__) @@ -13,12 +14,18 @@ logger = logging.getLogger(__name__) class Agents(AbstractResource): urls = ["/api/agents"] + def __init__(self, island_event_queue: IIslandEventQueue): + self._island_event_queue = island_event_queue + def post(self): try: # Just parse for now agent_registration_data = AgentRegistrationData(**request.json) logger.debug(f"Agent registered: {agent_registration_data}") + self._island_event_queue( + IslandEventTopic.AGENT_CONNECTED, agent_registration_data=agent_registration_data + ) return make_response({}, HTTPStatus.NO_CONTENT) except (TypeError, ValueError, json.JSONDecodeError) as err: diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py index 3ab90feb6..8bb61250b 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py @@ -1,8 +1,12 @@ from http import HTTPStatus +from unittest.mock import MagicMock from uuid import UUID +import pytest +from tests.common import StubDIContainer from tests.unit_tests.monkey_island.conftest import get_url_for_resource +from monkey_island.cc.event_queue import IIslandEventQueue from monkey_island.cc.resources import Agents AGENTS_URL = get_url_for_resource(Agents) @@ -17,6 +21,15 @@ AGENT_REGISTRATION_DICT = { } +@pytest.fixture +def flask_client(build_flask_client): + container = StubDIContainer() + container.register_instance(IIslandEventQueue, MagicMock(spec=IIslandEventQueue)) + + with build_flask_client(container) as flask_client: + yield flask_client + + def test_agent_registration(flask_client): print(AGENTS_URL) resp = flask_client.post( From 41dbb92eefb64cfb44e05891b9b9cbc4b421c5dd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 10:09:25 -0400 Subject: [PATCH 02/12] UT: Remove debug print statement --- .../tests/unit_tests/monkey_island/cc/resources/test_agents.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py index 8bb61250b..e44c36c09 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py +++ b/monkey/tests/unit_tests/monkey_island/cc/resources/test_agents.py @@ -31,7 +31,6 @@ def flask_client(build_flask_client): def test_agent_registration(flask_client): - print(AGENTS_URL) resp = flask_client.post( AGENTS_URL, json=AGENT_REGISTRATION_DICT, From c95c2ffdf94e01bee90a06addd2af936845b50e9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 09:58:12 -0400 Subject: [PATCH 03/12] Island: Add handle_agent_registration event handler - Add handle_agent_registration callable class - Add/Update machine to the repository --- .../cc/island_event_handlers/__init__.py | 1 + .../handle_agent_registration.py | 74 ++++++++++ .../test_handle_agent_registration.py | 128 ++++++++++++++++++ 3 files changed, 203 insertions(+) create mode 100644 monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py create mode 100644 monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py diff --git a/monkey/monkey_island/cc/island_event_handlers/__init__.py b/monkey/monkey_island/cc/island_event_handlers/__init__.py index 11343c2fa..bd390f447 100644 --- a/monkey/monkey_island/cc/island_event_handlers/__init__.py +++ b/monkey/monkey_island/cc/island_event_handlers/__init__.py @@ -1,3 +1,4 @@ +from .handle_agent_registration import handle_agent_registration from .reset_agent_configuration import reset_agent_configuration from .reset_machine_repository import reset_machine_repository from .set_agent_configuration_per_island_mode import set_agent_configuration_per_island_mode diff --git a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py new file mode 100644 index 000000000..9e3ead2b6 --- /dev/null +++ b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py @@ -0,0 +1,74 @@ +from contextlib import suppress +from typing import Optional + +from common import AgentRegistrationData +from monkey_island.cc.models import Machine +from monkey_island.cc.repository import IMachineRepository, UnknownRecordError + + +class handle_agent_registration: + """ + Update repositories when a new agent registers + """ + + def __init__(self, machine_repository: IMachineRepository): + self._machine_repository = machine_repository + + def __call__(self, agent_registration_data: AgentRegistrationData): + self._update_machine_repository(agent_registration_data) + + def _update_machine_repository(self, agent_registration_data: AgentRegistrationData): + machine = self._find_existing_machine_to_update(agent_registration_data) + + if machine is None: + machine = Machine(id=self._machine_repository.get_new_id()) + + self._upsert_machine(machine, agent_registration_data) + + def _find_existing_machine_to_update( + self, agent_registration_data: AgentRegistrationData + ) -> Optional[Machine]: + with suppress(UnknownRecordError): + return self._machine_repository.get_machine_by_hardware_id( + agent_registration_data.machine_hardware_id + ) + + for network_interface in agent_registration_data.network_interfaces: + with suppress(UnknownRecordError): + # NOTE: For now, assume IPs are unique. In reality, two machines could share the + # same IP if there's a router between them. + return self._machine_repository.get_machines_by_ip(network_interface.ip)[0] + + return None + + def _upsert_machine( + self, existing_machine: Machine, agent_registration_data: AgentRegistrationData + ): + updated_machine = existing_machine.copy() + + self._update_hardware_id(updated_machine, agent_registration_data) + self._update_network_interfaces(updated_machine, agent_registration_data) + + self._machine_repository.upsert_machine(updated_machine) + + def _update_hardware_id(self, machine: Machine, agent_registration_data: AgentRegistrationData): + if ( + machine.hardware_id is not None + and machine.hardware_id != agent_registration_data.machine_hardware_id + ): + raise Exception( + f"Hardware ID mismatch:\n\tMachine: {machine}\n\t" + f"AgentRegistrationData: {agent_registration_data}" + ) + + machine.hardware_id = agent_registration_data.machine_hardware_id + + def _update_network_interfaces( + self, machine: Machine, agent_registration_data: AgentRegistrationData + ): + updated_network_interfaces = set(machine.network_interfaces) + updated_network_interfaces = updated_network_interfaces.union( + agent_registration_data.network_interfaces + ) + + machine.network_interfaces = sorted(updated_network_interfaces) diff --git a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py new file mode 100644 index 000000000..37e95b14f --- /dev/null +++ b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py @@ -0,0 +1,128 @@ +from ipaddress import IPv4Address, IPv4Interface +from itertools import count +from typing import Sequence +from unittest.mock import MagicMock +from uuid import UUID + +import pytest + +from common import AgentRegistrationData +from monkey_island.cc.island_event_handlers import handle_agent_registration +from monkey_island.cc.models import Machine +from monkey_island.cc.repository import IMachineRepository, UnknownRecordError + +AGENT_ID = UUID("860aff5b-d2af-43ea-afb5-62bac3d30b7e") + +SEED_ID = 10 + +MACHINE = Machine( + id=2, + hardware_id=5, + network_interfaces=[IPv4Interface("192.168.2.2/24")], +) + +AGENT_REGISTRATION_DATA = AgentRegistrationData( + id=AGENT_ID, + machine_hardware_id=MACHINE.hardware_id, + start_time=0, + parent_id=None, + cc_server="192.168.1.1:5000", + network_interfaces=[IPv4Interface("192.168.1.2/24")], +) + + +@pytest.fixture +def machine_repository() -> IMachineRepository: + machine_repository = MagicMock(spec=IMachineRepository) + machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID)) + machine_repository.upsert_machine = MagicMock() + return machine_repository + + +@pytest.fixture +def handler(machine_repository) -> handle_agent_registration: + return handle_agent_registration(machine_repository) + + +def test_new_machine_added(handler, machine_repository): + expected_machine = Machine( + id=SEED_ID, + hardware_id=AGENT_REGISTRATION_DATA.machine_hardware_id, + network_interfaces=AGENT_REGISTRATION_DATA.network_interfaces, + ) + machine_repository.get_machine_by_hardware_id = MagicMock(side_effect=UnknownRecordError) + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) + handler(AGENT_REGISTRATION_DATA) + + machine_repository.upsert_machine.assert_called_once() + new_machine = machine_repository.upsert_machine.call_args_list[0][0][0] + + assert new_machine == expected_machine + + +def test_existing_machine_updated__hardware_id(handler, machine_repository): + expected_updated_machine = Machine( + id=MACHINE.id, + hardware_id=MACHINE.hardware_id, + network_interfaces=[ + AGENT_REGISTRATION_DATA.network_interfaces[0], + MACHINE.network_interfaces[0], + ], + ) + machine_repository.get_machine_by_hardware_id = MagicMock(return_value=MACHINE) + handler(AGENT_REGISTRATION_DATA) + + machine_repository.upsert_machine.assert_called_once() + machine_repository.upsert_machine.assert_called_with(expected_updated_machine) + + +def test_existing_machine_updated__find_by_ip(handler, machine_repository): + agent_registration_data = AgentRegistrationData( + id=AGENT_ID, + machine_hardware_id=5, + start_time=0, + parent_id=None, + cc_server="192.168.1.1:5000", + network_interfaces=[ + IPv4Interface("192.168.1.2/24"), + IPv4Interface("192.168.1.4/24"), + IPv4Interface("192.168.1.5/24"), + ], + ) + + existing_machine = Machine( + id=1, + network_interfaces=[agent_registration_data.network_interfaces[-1]], + ) + + def get_machines_by_ip(ip: IPv4Address) -> Sequence[Machine]: + if ip == existing_machine.network_interfaces[0].ip: + return [existing_machine] + + raise UnknownRecordError + + expected_updated_machine = existing_machine.copy() + expected_updated_machine.hardware_id = agent_registration_data.machine_hardware_id + expected_updated_machine.network_interfaces = agent_registration_data.network_interfaces + + machine_repository.get_machine_by_hardware_id = MagicMock(side_effect=UnknownRecordError) + machine_repository.get_machines_by_ip = MagicMock(side_effect=get_machines_by_ip) + + handler(agent_registration_data) + + machine_repository.upsert_machine.assert_called_once() + machine_repository.upsert_machine.assert_called_with(expected_updated_machine) + + +def test_hardware_id_mismatch(handler, machine_repository): + existing_machine = Machine( + id=1, + hardware_id=AGENT_REGISTRATION_DATA.machine_hardware_id + 99, + network_interfaces=AGENT_REGISTRATION_DATA.network_interfaces, + ) + + machine_repository.get_machine_by_hardware_id = MagicMock(side_effect=UnknownRecordError) + machine_repository.get_machines_by_ip = MagicMock(return_value=[existing_machine]) + + with pytest.raises(Exception): + handler(AGENT_REGISTRATION_DATA) From 1e8a60c8902bcb1f4cb85413625717856f2e7d88 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 10:29:55 -0400 Subject: [PATCH 04/12] Island: Add new agent to repository on agent registration --- .../handle_agent_registration.py | 24 +++++++++++---- .../test_handle_agent_registration.py | 30 ++++++++++++++++--- 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py index 9e3ead2b6..eaea53ffb 100644 --- a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py +++ b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py @@ -2,8 +2,8 @@ from contextlib import suppress from typing import Optional from common import AgentRegistrationData -from monkey_island.cc.models import Machine -from monkey_island.cc.repository import IMachineRepository, UnknownRecordError +from monkey_island.cc.models import Agent, Machine +from monkey_island.cc.repository import IAgentRepository, IMachineRepository, UnknownRecordError class handle_agent_registration: @@ -11,13 +11,15 @@ class handle_agent_registration: Update repositories when a new agent registers """ - def __init__(self, machine_repository: IMachineRepository): + def __init__(self, machine_repository: IMachineRepository, agent_repository: IAgentRepository): self._machine_repository = machine_repository + self._agent_repository = agent_repository def __call__(self, agent_registration_data: AgentRegistrationData): - self._update_machine_repository(agent_registration_data) + machine = self._update_machine_repository(agent_registration_data) + self._add_agent(agent_registration_data, machine) - def _update_machine_repository(self, agent_registration_data: AgentRegistrationData): + def _update_machine_repository(self, agent_registration_data: AgentRegistrationData) -> Machine: machine = self._find_existing_machine_to_update(agent_registration_data) if machine is None: @@ -25,6 +27,8 @@ class handle_agent_registration: self._upsert_machine(machine, agent_registration_data) + return machine + def _find_existing_machine_to_update( self, agent_registration_data: AgentRegistrationData ) -> Optional[Machine]: @@ -72,3 +76,13 @@ class handle_agent_registration: ) machine.network_interfaces = sorted(updated_network_interfaces) + + def _add_agent(self, agent_registration_data: AgentRegistrationData, machine: Machine): + new_agent = Agent( + id=agent_registration_data.id, + machine_id=machine.id, + start_time=agent_registration_data.start_time, + parent_id=agent_registration_data.parent_id, + cc_server=agent_registration_data.cc_server, + ) + self._agent_repository.upsert_agent(new_agent) diff --git a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py index 37e95b14f..91d392368 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py +++ b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py @@ -8,8 +8,8 @@ import pytest from common import AgentRegistrationData from monkey_island.cc.island_event_handlers import handle_agent_registration -from monkey_island.cc.models import Machine -from monkey_island.cc.repository import IMachineRepository, UnknownRecordError +from monkey_island.cc.models import Agent, Machine +from monkey_island.cc.repository import IAgentRepository, IMachineRepository, UnknownRecordError AGENT_ID = UUID("860aff5b-d2af-43ea-afb5-62bac3d30b7e") @@ -36,12 +36,21 @@ def machine_repository() -> IMachineRepository: machine_repository = MagicMock(spec=IMachineRepository) machine_repository.get_new_id = MagicMock(side_effect=count(SEED_ID)) machine_repository.upsert_machine = MagicMock() + machine_repository.get_machine_by_hardware_id = MagicMock(side_effect=UnknownRecordError) + machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) return machine_repository @pytest.fixture -def handler(machine_repository) -> handle_agent_registration: - return handle_agent_registration(machine_repository) +def agent_repository() -> IAgentRepository: + agent_repository = MagicMock(spec=IAgentRepository) + agent_repository.upsert_agent = MagicMock() + return agent_repository + + +@pytest.fixture +def handler(machine_repository, agent_repository) -> handle_agent_registration: + return handle_agent_registration(machine_repository, agent_repository) def test_new_machine_added(handler, machine_repository): @@ -126,3 +135,16 @@ def test_hardware_id_mismatch(handler, machine_repository): with pytest.raises(Exception): handler(AGENT_REGISTRATION_DATA) + + +def test_add_agent(handler, agent_repository): + expected_agent = Agent( + id=AGENT_REGISTRATION_DATA.id, + machine_id=SEED_ID, + start_time=AGENT_REGISTRATION_DATA.start_time, + parent_id=AGENT_REGISTRATION_DATA.parent_id, + cc_server=AGENT_REGISTRATION_DATA.cc_server, + ) + handler(AGENT_REGISTRATION_DATA) + + agent_repository.upsert_agent.assert_called_with(expected_agent) From c0870e6696b6ad35f9c39d6fe94eec1cbea58270 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 12:04:37 -0400 Subject: [PATCH 05/12] Island: Add new node communication on agent registration --- .../handle_agent_registration.py | 42 ++++++++- .../test_handle_agent_registration.py | 87 +++++++++++++++---- 2 files changed, 109 insertions(+), 20 deletions(-) diff --git a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py index eaea53ffb..f6add225f 100644 --- a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py +++ b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py @@ -1,9 +1,16 @@ from contextlib import suppress +from ipaddress import IPv4Address, IPv4Interface from typing import Optional from common import AgentRegistrationData -from monkey_island.cc.models import Agent, Machine -from monkey_island.cc.repository import IAgentRepository, IMachineRepository, UnknownRecordError +from common.network.network_utils import address_to_ip_port +from monkey_island.cc.models import Agent, CommunicationType, Machine +from monkey_island.cc.repository import ( + IAgentRepository, + IMachineRepository, + INodeRepository, + UnknownRecordError, +) class handle_agent_registration: @@ -11,13 +18,20 @@ class handle_agent_registration: Update repositories when a new agent registers """ - def __init__(self, machine_repository: IMachineRepository, agent_repository: IAgentRepository): + def __init__( + self, + machine_repository: IMachineRepository, + agent_repository: IAgentRepository, + node_repository: INodeRepository, + ): self._machine_repository = machine_repository self._agent_repository = agent_repository + self._node_repository = node_repository def __call__(self, agent_registration_data: AgentRegistrationData): machine = self._update_machine_repository(agent_registration_data) self._add_agent(agent_registration_data, machine) + self._add_node_communication(agent_registration_data, machine) def _update_machine_repository(self, agent_registration_data: AgentRegistrationData) -> Machine: machine = self._find_existing_machine_to_update(agent_registration_data) @@ -86,3 +100,25 @@ class handle_agent_registration: cc_server=agent_registration_data.cc_server, ) self._agent_repository.upsert_agent(new_agent) + + def _add_node_communication( + self, agent_registration_data: AgentRegistrationData, src_machine: Machine + ): + dst_machine = self._get_or_create_cc_machine(agent_registration_data.cc_server) + + self._node_repository.upsert_communication( + src_machine.id, dst_machine.id, CommunicationType.CC + ) + + def _get_or_create_cc_machine(self, cc_server: str) -> Machine: + dst_ip = IPv4Address(address_to_ip_port(cc_server)[0]) + + try: + return self._machine_repository.get_machines_by_ip(dst_ip)[0] + except UnknownRecordError: + new_machine = Machine( + id=self._machine_repository.get_new_id(), network_interfaces=[IPv4Interface(dst_ip)] + ) + self._machine_repository.upsert_machine(new_machine) + + return new_machine diff --git a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py index 91d392368..1ebfb59f7 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py +++ b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py @@ -8,8 +8,13 @@ import pytest from common import AgentRegistrationData from monkey_island.cc.island_event_handlers import handle_agent_registration -from monkey_island.cc.models import Agent, Machine -from monkey_island.cc.repository import IAgentRepository, IMachineRepository, UnknownRecordError +from monkey_island.cc.models import Agent, CommunicationType, Machine +from monkey_island.cc.repository import ( + IAgentRepository, + IMachineRepository, + INodeRepository, + UnknownRecordError, +) AGENT_ID = UUID("860aff5b-d2af-43ea-afb5-62bac3d30b7e") @@ -49,8 +54,25 @@ def agent_repository() -> IAgentRepository: @pytest.fixture -def handler(machine_repository, agent_repository) -> handle_agent_registration: - return handle_agent_registration(machine_repository, agent_repository) +def node_repository() -> INodeRepository: + node_repository = MagicMock(spec=INodeRepository) + node_repository.upsert_communication = MagicMock() + return node_repository + + +@pytest.fixture +def handler(machine_repository, agent_repository, node_repository) -> handle_agent_registration: + return handle_agent_registration(machine_repository, agent_repository, node_repository) + + +def build_get_machines_by_ip(ip_to_match: IPv4Address, machine_to_return: Machine): + def get_machines_by_ip(ip: IPv4Address) -> Sequence[Machine]: + if ip == ip_to_match: + return [machine_to_return] + + raise UnknownRecordError + + return get_machines_by_ip def test_new_machine_added(handler, machine_repository): @@ -61,12 +83,10 @@ def test_new_machine_added(handler, machine_repository): ) machine_repository.get_machine_by_hardware_id = MagicMock(side_effect=UnknownRecordError) machine_repository.get_machines_by_ip = MagicMock(side_effect=UnknownRecordError) + handler(AGENT_REGISTRATION_DATA) - machine_repository.upsert_machine.assert_called_once() - new_machine = machine_repository.upsert_machine.call_args_list[0][0][0] - - assert new_machine == expected_machine + machine_repository.upsert_machine.assert_any_call(expected_machine) def test_existing_machine_updated__hardware_id(handler, machine_repository): @@ -79,10 +99,10 @@ def test_existing_machine_updated__hardware_id(handler, machine_repository): ], ) machine_repository.get_machine_by_hardware_id = MagicMock(return_value=MACHINE) + handler(AGENT_REGISTRATION_DATA) - machine_repository.upsert_machine.assert_called_once() - machine_repository.upsert_machine.assert_called_with(expected_updated_machine) + machine_repository.upsert_machine.assert_any_call(expected_updated_machine) def test_existing_machine_updated__find_by_ip(handler, machine_repository): @@ -104,11 +124,9 @@ def test_existing_machine_updated__find_by_ip(handler, machine_repository): network_interfaces=[agent_registration_data.network_interfaces[-1]], ) - def get_machines_by_ip(ip: IPv4Address) -> Sequence[Machine]: - if ip == existing_machine.network_interfaces[0].ip: - return [existing_machine] - - raise UnknownRecordError + get_machines_by_ip = build_get_machines_by_ip( + existing_machine.network_interfaces[0].ip, existing_machine + ) expected_updated_machine = existing_machine.copy() expected_updated_machine.hardware_id = agent_registration_data.machine_hardware_id @@ -119,8 +137,7 @@ def test_existing_machine_updated__find_by_ip(handler, machine_repository): handler(agent_registration_data) - machine_repository.upsert_machine.assert_called_once() - machine_repository.upsert_machine.assert_called_with(expected_updated_machine) + machine_repository.upsert_machine.assert_any_call(expected_updated_machine) def test_hardware_id_mismatch(handler, machine_repository): @@ -148,3 +165,39 @@ def test_add_agent(handler, agent_repository): handler(AGENT_REGISTRATION_DATA) agent_repository.upsert_agent.assert_called_with(expected_agent) + + +def test_add_node_connection(handler, machine_repository, node_repository): + island_machine = Machine( + id=1, + hardware_id=99, + island=True, + network_interfaces=[IPv4Interface("192.168.1.1/24")], + ) + get_machines_by_ip = build_get_machines_by_ip( + island_machine.network_interfaces[0].ip, island_machine + ) + machine_repository.get_machines_by_ip = MagicMock(side_effect=get_machines_by_ip) + machine_repository.get_machine_by_hardware_id = MagicMock(return_value=MACHINE) + + handler(AGENT_REGISTRATION_DATA) + + node_repository.upsert_communication.assert_called_once() + node_repository.upsert_communication.assert_called_with( + MACHINE.id, island_machine.id, CommunicationType.CC + ) + + +def test_add_node_connection__unknown_server(handler, machine_repository, node_repository): + expected_new_server_machine = Machine( + id=SEED_ID, + network_interfaces=[IPv4Interface("192.168.1.1/32")], + ) + + machine_repository.get_machine_by_hardware_id = MagicMock(return_value=MACHINE) + handler(AGENT_REGISTRATION_DATA) + + machine_repository.upsert_machine.assert_called_with(expected_new_server_machine) + node_repository.upsert_communication.assert_called_with( + MACHINE.id, SEED_ID, CommunicationType.CC + ) From 4101f7509a5530585c1625d65515e598e67d9c3a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 12:35:13 -0400 Subject: [PATCH 06/12] Island: Handle IP address collisions when updating Machine on agent reg --- .../handle_agent_registration.py | 18 ++++++++--- .../test_handle_agent_registration.py | 32 +++++++++++++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py index f6add225f..c9bf5f611 100644 --- a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py +++ b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py @@ -1,6 +1,6 @@ from contextlib import suppress from ipaddress import IPv4Address, IPv4Interface -from typing import Optional +from typing import List, Optional from common import AgentRegistrationData from common.network.network_utils import address_to_ip_port @@ -84,11 +84,21 @@ class handle_agent_registration: def _update_network_interfaces( self, machine: Machine, agent_registration_data: AgentRegistrationData ): - updated_network_interfaces = set(machine.network_interfaces) - updated_network_interfaces = updated_network_interfaces.union( - agent_registration_data.network_interfaces + updated_network_interfaces: List[IPv4Interface] = [] + agent_registration_data_ips = set( + map(lambda iface: iface.ip, agent_registration_data.network_interfaces) ) + # Prefer interfaces provided by the AgentRegistrationData to those in the Machine record. + # The AgentRegistrationData was collected while running on the machine, whereas the Machine + # data may have only been collected from a scan. For example, the Machine and + # AgentRedistrationData may have the same IP with a different subnet mask. + for interface in machine.network_interfaces: + if interface.ip not in agent_registration_data_ips: + updated_network_interfaces.append(interface) + + updated_network_interfaces.extend(agent_registration_data.network_interfaces) + machine.network_interfaces = sorted(updated_network_interfaces) def _add_agent(self, agent_registration_data: AgentRegistrationData, machine: Machine): diff --git a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py index 1ebfb59f7..4386469be 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py +++ b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py @@ -201,3 +201,35 @@ def test_add_node_connection__unknown_server(handler, machine_repository, node_r node_repository.upsert_communication.assert_called_with( MACHINE.id, SEED_ID, CommunicationType.CC ) + + +def test_machine_interfaces_updated(handler, machine_repository): + existing_machine = Machine( + id=SEED_ID, + hardware_id=AGENT_REGISTRATION_DATA.machine_hardware_id, + network_interfaces=[IPv4Interface("192.168.1.2/32"), IPv4Interface("192.168.1.5/32")], + ) + machine_repository.get_machine_by_hardware_id = MagicMock(return_value=existing_machine) + agent_registration_data = AgentRegistrationData( + id=AGENT_ID, + machine_hardware_id=MACHINE.hardware_id, + start_time=0, + parent_id=None, + cc_server="192.168.1.1:5000", + network_interfaces=[ + IPv4Interface("192.168.1.2/24"), + IPv4Interface("192.168.1.3/16"), + IPv4Interface("192.168.1.4/24"), + ], + ) + expected_network_interfaces = tuple( + sorted( + (*agent_registration_data.network_interfaces, existing_machine.network_interfaces[-1]) + ) + ) + + handler(agent_registration_data) + updated_machine = machine_repository.upsert_machine.call_args_list[0][0][0] + actual_network_interfaces = updated_machine.network_interfaces + + assert actual_network_interfaces == expected_network_interfaces From 271475570de09b43b4127a46bd50c8879a51f505 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 12:36:32 -0400 Subject: [PATCH 07/12] Island: Modify machine with AgentRegistrationData Since the Machine is passed to other methods, it should have the most up to date data. --- .../handle_agent_registration.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py index c9bf5f611..363335c5f 100644 --- a/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py +++ b/monkey/monkey_island/cc/island_event_handlers/handle_agent_registration.py @@ -59,15 +59,11 @@ class handle_agent_registration: return None - def _upsert_machine( - self, existing_machine: Machine, agent_registration_data: AgentRegistrationData - ): - updated_machine = existing_machine.copy() + def _upsert_machine(self, machine: Machine, agent_registration_data: AgentRegistrationData): + self._update_hardware_id(machine, agent_registration_data) + self._update_network_interfaces(machine, agent_registration_data) - self._update_hardware_id(updated_machine, agent_registration_data) - self._update_network_interfaces(updated_machine, agent_registration_data) - - self._machine_repository.upsert_machine(updated_machine) + self._machine_repository.upsert_machine(machine) def _update_hardware_id(self, machine: Machine, agent_registration_data: AgentRegistrationData): if ( From 71da0a559ac6ceda09ae3588c0dd998053e4fa6f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 14:56:06 -0400 Subject: [PATCH 08/12] Island: Register handle_agent_registration for AGENT_CONNECTED events --- monkey/monkey_island/cc/setup/island_event_handlers.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index ee37568f1..fa036a8fe 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -3,6 +3,7 @@ from functools import partial from common import DIContainer from monkey_island.cc.event_queue import IIslandEventQueue, IslandEventTopic from monkey_island.cc.island_event_handlers import ( + handle_agent_registration, reset_agent_configuration, reset_machine_repository, set_agent_configuration_per_island_mode, @@ -20,11 +21,20 @@ from monkey_island.cc.services.database import Database def setup_island_event_handlers(container: DIContainer): island_event_queue = container.resolve(IIslandEventQueue) + _subscribe_agent_registration_events(island_event_queue, container) _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) +def _subscribe_agent_registration_events( + island_event_queue: IIslandEventQueue, container: DIContainer +): + topic = IslandEventTopic.AGENT_CONNECTED + + island_event_queue.subscribe(topic, container.resolve(handle_agent_registration)) + + def _subscribe_reset_agent_configuration_events( island_event_queue: IIslandEventQueue, container: DIContainer ): From 371e495f171b1f773d2011284609a8984c64b95d Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 15:41:48 -0400 Subject: [PATCH 09/12] Island: Fix call to IIslandEventQueue.publish() --- monkey/monkey_island/cc/resources/agents.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/resources/agents.py b/monkey/monkey_island/cc/resources/agents.py index 4fb9837ff..9d6f701ec 100644 --- a/monkey/monkey_island/cc/resources/agents.py +++ b/monkey/monkey_island/cc/resources/agents.py @@ -23,8 +23,8 @@ class Agents(AbstractResource): agent_registration_data = AgentRegistrationData(**request.json) logger.debug(f"Agent registered: {agent_registration_data}") - self._island_event_queue( - IslandEventTopic.AGENT_CONNECTED, agent_registration_data=agent_registration_data + self._island_event_queue.publish( + IslandEventTopic.AGENT_REGISTERED, agent_registration_data=agent_registration_data ) return make_response({}, HTTPStatus.NO_CONTENT) From bed684d7d03295b6d021351b5410d6a7e0675fa1 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 15:42:28 -0400 Subject: [PATCH 10/12] Island: Rename IslandEventTopic.AGENT_{CONNECTED,REGISTERED} --- .../cc/event_queue/i_island_event_queue.py | 2 +- monkey/monkey_island/cc/setup/island_event_handlers.py | 2 +- .../cc/event_queue/test_pypubsub_island_event_queue.py | 10 +++++----- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py index bf1fcf2cc..cf123c3e2 100644 --- a/monkey/monkey_island/cc/event_queue/i_island_event_queue.py +++ b/monkey/monkey_island/cc/event_queue/i_island_event_queue.py @@ -5,7 +5,7 @@ from . import IslandEventSubscriber class IslandEventTopic(Enum): - AGENT_CONNECTED = auto() + AGENT_REGISTERED = auto() CLEAR_SIMULATION_DATA = auto() RESET_AGENT_CONFIGURATION = auto() SET_ISLAND_MODE = auto() diff --git a/monkey/monkey_island/cc/setup/island_event_handlers.py b/monkey/monkey_island/cc/setup/island_event_handlers.py index fa036a8fe..6b4cb8c53 100644 --- a/monkey/monkey_island/cc/setup/island_event_handlers.py +++ b/monkey/monkey_island/cc/setup/island_event_handlers.py @@ -30,7 +30,7 @@ def setup_island_event_handlers(container: DIContainer): def _subscribe_agent_registration_events( island_event_queue: IIslandEventQueue, container: DIContainer ): - topic = IslandEventTopic.AGENT_CONNECTED + topic = IslandEventTopic.AGENT_REGISTERED island_event_queue.subscribe(topic, container.resolve(handle_agent_registration)) diff --git a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py index 206009727..91af47fa0 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py +++ b/monkey/tests/unit_tests/monkey_island/cc/event_queue/test_pypubsub_island_event_queue.py @@ -40,7 +40,7 @@ def test_subscribe_publish__no_event_body( topic=IslandEventTopic.CLEAR_SIMULATION_DATA, subscriber=event_queue_subscriber ) - event_queue.publish(topic=IslandEventTopic.AGENT_CONNECTED) + event_queue.publish(topic=IslandEventTopic.AGENT_REGISTERED) event_queue.publish(topic=IslandEventTopic.CLEAR_SIMULATION_DATA) event_queue.publish(topic=IslandEventTopic.RESET_AGENT_CONFIGURATION) @@ -64,9 +64,9 @@ def test_subscribe_publish__with_event_body( event = "my event!" my_callable = MyCallable() - event_queue.subscribe(topic=IslandEventTopic.AGENT_CONNECTED, subscriber=my_callable) + event_queue.subscribe(topic=IslandEventTopic.AGENT_REGISTERED, subscriber=my_callable) - event_queue.publish(topic=IslandEventTopic.AGENT_CONNECTED, event=event) + event_queue.publish(topic=IslandEventTopic.AGENT_REGISTERED, event=event) event_queue.publish(topic=IslandEventTopic.CLEAR_SIMULATION_DATA) event_queue.publish(topic=IslandEventTopic.RESET_AGENT_CONFIGURATION) @@ -84,10 +84,10 @@ def test_keep_subscriber_in_scope(event_queue: IIslandEventQueue): def subscribe(): # fn will go out of scope after subscribe() returns. fn = MyCallable() - event_queue.subscribe(topic=IslandEventTopic.AGENT_CONNECTED, subscriber=fn) + event_queue.subscribe(topic=IslandEventTopic.AGENT_REGISTERED, subscriber=fn) subscribe() - event_queue.publish(topic=IslandEventTopic.AGENT_CONNECTED) + event_queue.publish(topic=IslandEventTopic.AGENT_REGISTERED) assert MyCallable.called From 5b61767575ebe1f415fcb3f5ce58e448eda8f700 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 20:34:30 -0400 Subject: [PATCH 11/12] Project: Remove IslandEventTopics from vulture_allowlist --- vulture_allowlist.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/vulture_allowlist.py b/vulture_allowlist.py index 40a2f5cdb..156db5df5 100644 --- a/vulture_allowlist.py +++ b/vulture_allowlist.py @@ -318,7 +318,3 @@ SCANNED EXPLOITED CC CC_TUNNEL - -IslandEventTopic.AGENT_CONNECTED -IslandEventTopic.CLEAR_SIMULATION_DATA -IslandEventTopic.RESET_AGENT_CONFIGURATION From 284f58afaffbe906995c22fc338be95a26406113 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 21 Sep 2022 20:37:30 -0400 Subject: [PATCH 12/12] UT: Sort updated_machine.network_interfaces --- .../test_handle_agent_registration.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py index 4386469be..c6070f1ed 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py +++ b/monkey/tests/unit_tests/monkey_island/cc/island_event_handlers/test_handle_agent_registration.py @@ -222,14 +222,12 @@ def test_machine_interfaces_updated(handler, machine_repository): IPv4Interface("192.168.1.4/24"), ], ) - expected_network_interfaces = tuple( - sorted( - (*agent_registration_data.network_interfaces, existing_machine.network_interfaces[-1]) - ) + expected_network_interfaces = sorted( + (*agent_registration_data.network_interfaces, existing_machine.network_interfaces[-1]) ) handler(agent_registration_data) updated_machine = machine_repository.upsert_machine.call_args_list[0][0][0] - actual_network_interfaces = updated_machine.network_interfaces + actual_network_interfaces = sorted(updated_machine.network_interfaces) assert actual_network_interfaces == expected_network_interfaces