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