forked from p34709852/monkey
Island: Handle IP address collisions when updating Machine on agent reg
This commit is contained in:
parent
c0870e6696
commit
4101f7509a
|
@ -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):
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in New Issue