Island: Handle IP address collisions when updating Machine on agent reg

This commit is contained in:
Mike Salvatore 2022-09-21 12:35:13 -04:00
parent c0870e6696
commit 4101f7509a
2 changed files with 46 additions and 4 deletions

View File

@ -1,6 +1,6 @@
from contextlib import suppress from contextlib import suppress
from ipaddress import IPv4Address, IPv4Interface from ipaddress import IPv4Address, IPv4Interface
from typing import Optional from typing import List, Optional
from common import AgentRegistrationData from common import AgentRegistrationData
from common.network.network_utils import address_to_ip_port from common.network.network_utils import address_to_ip_port
@ -84,11 +84,21 @@ class handle_agent_registration:
def _update_network_interfaces( def _update_network_interfaces(
self, machine: Machine, agent_registration_data: AgentRegistrationData self, machine: Machine, agent_registration_data: AgentRegistrationData
): ):
updated_network_interfaces = set(machine.network_interfaces) updated_network_interfaces: List[IPv4Interface] = []
updated_network_interfaces = updated_network_interfaces.union( agent_registration_data_ips = set(
agent_registration_data.network_interfaces 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) machine.network_interfaces = sorted(updated_network_interfaces)
def _add_agent(self, agent_registration_data: AgentRegistrationData, machine: Machine): def _add_agent(self, agent_registration_data: AgentRegistrationData, machine: Machine):

View File

@ -201,3 +201,35 @@ def test_add_node_connection__unknown_server(handler, machine_repository, node_r
node_repository.upsert_communication.assert_called_with( node_repository.upsert_communication.assert_called_with(
MACHINE.id, SEED_ID, CommunicationType.CC 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