From 6df2c29e309704a88d31b2c9ac3fbf213344aace Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 17:13:27 +0530 Subject: [PATCH 01/12] Common: Update cc_server's type in AgentRegistrationData to SocketAddress --- monkey/common/agent_registration_data.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/common/agent_registration_data.py b/monkey/common/agent_registration_data.py index 6d45df89f..5d281be25 100644 --- a/monkey/common/agent_registration_data.py +++ b/monkey/common/agent_registration_data.py @@ -7,7 +7,7 @@ from pydantic import validator from .base_models import InfectionMonkeyBaseModel from .transforms import make_immutable_sequence -from .types import HardwareID +from .types import HardwareID, SocketAddress class AgentRegistrationData(InfectionMonkeyBaseModel): @@ -15,7 +15,7 @@ class AgentRegistrationData(InfectionMonkeyBaseModel): machine_hardware_id: HardwareID start_time: datetime parent_id: Optional[UUID] - cc_server: str + cc_server: SocketAddress network_interfaces: Sequence[IPv4Interface] _make_immutable_sequence = validator("network_interfaces", pre=True, allow_reuse=True)( From ac633a6e75850eb80fda6497d1efbd3fc5e46611 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 17:16:17 +0530 Subject: [PATCH 02/12] Agent: Change argument in InfectionMonkey._register_agent() to type SocketAddress --- monkey/infection_monkey/monkey.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e86bdbee1..071f742b1 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -131,7 +131,7 @@ class InfectionMonkey: server_address=server, island_api_client=self._island_api_client ) self._control_channel = ControlChannel(server, get_agent_id(), self._island_api_client) - self._register_agent(server) + self._register_agent(self._island_address) # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object @@ -180,7 +180,7 @@ class InfectionMonkey: return server, island_api_client - def _register_agent(self, server: str): + def _register_agent(self, server: SocketAddress): agent_registration_data = AgentRegistrationData( id=get_agent_id(), machine_hardware_id=get_machine_id(), From 284ec3d1197fa250c31d54d91b94cc5812fe322a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 17:46:06 +0530 Subject: [PATCH 03/12] Agent: Use SocketAddress type in handle_agent_registration._get_or_create_cc_machine() --- .../cc/island_event_handlers/handle_agent_registration.py | 6 +++--- 1 file changed, 3 insertions(+), 3 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 363335c5f..2994f03c3 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 @@ -3,7 +3,7 @@ from ipaddress import IPv4Address, IPv4Interface from typing import List, Optional from common import AgentRegistrationData -from common.network.network_utils import address_to_ip_port +from common.types import SocketAddress from monkey_island.cc.models import Agent, CommunicationType, Machine from monkey_island.cc.repository import ( IAgentRepository, @@ -116,8 +116,8 @@ class handle_agent_registration: 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]) + def _get_or_create_cc_machine(self, cc_server: SocketAddress) -> Machine: + dst_ip = cc_server.ip try: return self._machine_repository.get_machines_by_ip(dst_ip)[0] From 6a29702846746e98c08e4c349d014c168a3f43d3 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 17:54:03 +0530 Subject: [PATCH 04/12] Island: Modify cc_server type in Agent pydantic model --- monkey/monkey_island/cc/models/agent.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monkey/monkey_island/cc/models/agent.py b/monkey/monkey_island/cc/models/agent.py index c5ae10d22..3fb384e86 100644 --- a/monkey/monkey_island/cc/models/agent.py +++ b/monkey/monkey_island/cc/models/agent.py @@ -4,6 +4,7 @@ from typing import Optional from pydantic import Field from common.base_models import MutableInfectionMonkeyBaseModel +from common.types import SocketAddress from . import AgentID, MachineID @@ -26,7 +27,7 @@ class Agent(MutableInfectionMonkeyBaseModel): parent_id: Optional[AgentID] = Field(allow_mutation=False) """The ID of the parent agent that spawned this agent""" - cc_server: str = Field(default="") + cc_server: Optional[SocketAddress] """The address that the agent used to communicate with the island""" log_contents: str = Field(default="") From 8a609e0871c3768374c047c522692c718a061153 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 18:13:48 +0530 Subject: [PATCH 05/12] Agent: Simplify logic in SocketAddress object creation in monkey.py and network/relay/utils.py --- monkey/infection_monkey/monkey.py | 6 ++---- monkey/infection_monkey/network/relay/utils.py | 4 +--- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 071f742b1..a1cbdc2f2 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -3,7 +3,7 @@ import logging import os import subprocess import sys -from ipaddress import IPv4Address, IPv4Interface +from ipaddress import IPv4Interface from pathlib import Path, WindowsPath from typing import List, Mapping, Optional, Tuple @@ -123,9 +123,7 @@ class InfectionMonkey: self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._cmd_island_port = int(self._cmd_island_port) - self._island_address = SocketAddress( - IPv4Address(self._cmd_island_ip), self._cmd_island_port - ) + self._island_address = SocketAddress(self._cmd_island_ip, self._cmd_island_port) self._control_client = ControlClient( server_address=server, island_api_client=self._island_api_client diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 3eca4ac01..737a07c9a 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,7 +1,6 @@ import logging import socket from contextlib import suppress -from ipaddress import IPv4Address from typing import Dict, Iterable, Iterator, Optional from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT @@ -90,8 +89,7 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): def _send_remove_from_waitlist_control_message_to_relay(server: str): - ip, port = address_to_ip_port(server) - server_address = SocketAddress(IPv4Address(ip), int(port)) + server_address = SocketAddress(*address_to_ip_port(server)) notify_disconnect(server_address) From 0ae653fb729cf5d0028b13666ca0d2b5674a0ba0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 18:19:49 +0530 Subject: [PATCH 06/12] Agent: Let InfectionMonkey._cmd_island_port stay a str No part of the code in InfectionMonkey requires _cmd_island_port to be an int now --- monkey/infection_monkey/monkey.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index a1cbdc2f2..80f6d9735 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -119,9 +119,7 @@ class InfectionMonkey: self._agent_event_serializer_registry = self._setup_agent_event_serializers() server, self._island_api_client = self._connect_to_island_api() - # TODO: `address_to_port()` should return the port as an integer. self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) - self._cmd_island_port = int(self._cmd_island_port) self._island_address = SocketAddress(self._cmd_island_ip, self._cmd_island_port) From 3b192a869e45c1b0e77cc075a2e670d5a3a89752 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 18:36:08 +0530 Subject: [PATCH 07/12] UT: Replace test data with SocketAddress wherever required --- .../island_api_client/test_http_island_api_client.py | 3 ++- .../test_handle_agent_registration.py | 7 ++++--- .../cc/models/test_agent_registration_message.py | 8 +++++--- .../cc/repository/test_mongo_agent_repository.py | 4 +++- .../unit_tests/monkey_island/cc/resources/test_agents.py | 2 +- 5 files changed, 15 insertions(+), 9 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py index a861a4849..56a480082 100644 --- a/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py +++ b/monkey/tests/unit_tests/infection_monkey/island_api_client/test_http_island_api_client.py @@ -11,6 +11,7 @@ from common.agent_event_serializers import ( ) from common.agent_events import AbstractAgentEvent from common.agent_registration_data import AgentRegistrationData +from common.types import SocketAddress from infection_monkey.island_api_client import ( HTTPIslandAPIClient, IslandAPIConnectionError, @@ -20,7 +21,7 @@ from infection_monkey.island_api_client import ( IslandAPITimeoutError, ) -SERVER = "1.1.1.1:9999" +SERVER = SocketAddress(ip="1.1.1.1", port="9999") PBA_FILE = "dummy.pba" WINDOWS = "windows" AGENT_ID = UUID("80988359-a1cd-42a2-9b47-5b94b37cd673") 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 c6070f1ed..61337b022 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 @@ -7,6 +7,7 @@ from uuid import UUID import pytest from common import AgentRegistrationData +from common.types import SocketAddress from monkey_island.cc.island_event_handlers import handle_agent_registration from monkey_island.cc.models import Agent, CommunicationType, Machine from monkey_island.cc.repository import ( @@ -31,7 +32,7 @@ AGENT_REGISTRATION_DATA = AgentRegistrationData( machine_hardware_id=MACHINE.hardware_id, start_time=0, parent_id=None, - cc_server="192.168.1.1:5000", + cc_server=SocketAddress(ip="192.168.1.1", port="5000"), network_interfaces=[IPv4Interface("192.168.1.2/24")], ) @@ -111,7 +112,7 @@ def test_existing_machine_updated__find_by_ip(handler, machine_repository): machine_hardware_id=5, start_time=0, parent_id=None, - cc_server="192.168.1.1:5000", + cc_server=SocketAddress(ip="192.168.1.1", port="5000"), network_interfaces=[ IPv4Interface("192.168.1.2/24"), IPv4Interface("192.168.1.4/24"), @@ -215,7 +216,7 @@ def test_machine_interfaces_updated(handler, machine_repository): machine_hardware_id=MACHINE.hardware_id, start_time=0, parent_id=None, - cc_server="192.168.1.1:5000", + cc_server=SocketAddress(ip="192.168.1.1", port="5000"), network_interfaces=[ IPv4Interface("192.168.1.2/24"), IPv4Interface("192.168.1.3/16"), diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py index 2b755038c..d03761628 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py @@ -6,16 +6,18 @@ from uuid import UUID import pytest from common import AgentRegistrationData +from common.types import SocketAddress AGENT_ID = UUID("012e7238-7b81-4108-8c7f-0787bc3f3c10") PARENT_ID = UUID("0fc9afcb-1902-436b-bd5c-1ad194252484") +SOCKET_ADDRESS = SocketAddress(ip="192.168.1.1", port="5000") AGENT_REGISTRATION_MESSAGE_OBJECT_DICT = { "id": AGENT_ID, "machine_hardware_id": 2, "start_time": datetime.fromtimestamp(1660848408, tz=timezone.utc), "parent_id": PARENT_ID, - "cc_server": "192.168.1.1:5000", + "cc_server": SOCKET_ADDRESS, "network_interfaces": [IPv4Interface("10.0.0.1/24"), IPv4Interface("192.168.5.32/16")], } @@ -24,7 +26,7 @@ AGENT_REGISTRATION_MESSAGE_SIMPLE_DICT = { "machine_hardware_id": 2, "start_time": "2022-08-18T18:46:48+00:00", "parent_id": str(PARENT_ID), - "cc_server": "192.168.1.1:5000", + "cc_server": SOCKET_ADDRESS.dict(simplify=True), "network_interfaces": ["10.0.0.1/24", "192.168.5.32/16"], } @@ -85,7 +87,7 @@ def test_construct_invalid_field__value_error(key, value): ("machine_hardware_id", 99), ("start_time", 0), ("parent_id", AGENT_ID), - ("cc_server", "10.0.0.1:4999"), + ("cc_server", SOCKET_ADDRESS), ("network_interfaces", ["10.0.0.1/24"]), ], ) diff --git a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py index 938ea7412..aa804cd13 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py +++ b/monkey/tests/unit_tests/monkey_island/cc/repository/test_mongo_agent_repository.py @@ -6,6 +6,7 @@ from uuid import uuid4 import mongomock import pytest +from common.types import SocketAddress from monkey_island.cc.models import Agent from monkey_island.cc.repository import ( IAgentRepository, @@ -54,6 +55,7 @@ AGENTS = ( *RUNNING_AGENTS, *STOPPED_AGENTS, ) +CC_SERVER = SocketAddress(ip="127.0.0.1", port="1984") @pytest.fixture @@ -116,7 +118,7 @@ def test_upsert_agent__insert_empty_repository(empty_agent_repository): def test_upsert_agent__update(agent_repository): agents = deepcopy(AGENTS) agents[0].stop_time = datetime.now() - agents[0].cc_server = "127.0.0.1:1984" + agents[0].cc_server = CC_SERVER agent_repository.upsert_agent(agents[0]) 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 e44c36c09..97a9163de 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 @@ -16,7 +16,7 @@ AGENT_REGISTRATION_DICT = { "machine_hardware_id": 1, "start_time": 0, "parent_id": UUID("9d55ba33-95c2-417d-bd86-d3d11e47daeb"), - "cc_server": "10.0.0.1:5000", + "cc_server": {"ip": "10.0.0.1", "port": "5000"}, "network_interfaces": ["10.1.1.2/24"], } From 8b0ebfc3a7db87e069810f2e93e43e4069a27fc9 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 18:49:22 +0530 Subject: [PATCH 08/12] UT: Add/modify tests as per SocketAddress related changes --- .../unit_tests/monkey_island/cc/models/test_agent.py | 11 ++++++----- .../cc/models/test_agent_registration_message.py | 3 ++- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py index fb528fd51..925db210a 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent.py @@ -27,7 +27,7 @@ def test_constructor__defaults_from_objects(): a = Agent(**AGENT_OBJECT_DICT) assert a.stop_time is None - assert a.cc_server == "" + assert a.cc_server is None def test_constructor__defaults_from_simple_dict(): @@ -37,7 +37,7 @@ def test_constructor__defaults_from_simple_dict(): assert a.parent_id is None assert a.stop_time is None - assert a.cc_server == "" + assert a.cc_server is None assert a.log_contents == "" @@ -45,7 +45,7 @@ def test_to_dict(): a = Agent(**AGENT_OBJECT_DICT) agent_simple_dict = AGENT_SIMPLE_DICT.copy() agent_simple_dict["stop_time"] = None - agent_simple_dict["cc_server"] = "" + agent_simple_dict["cc_server"] = None agent_simple_dict["log_contents"] = "" assert a.dict(simplify=True) == agent_simple_dict @@ -59,7 +59,7 @@ def test_to_dict(): ("start_time", None), ("stop_time", []), ("parent_id", 2.1), - ("cc_server", []), + ("cc_server", [1]), ("log_contents", None), ], ) @@ -77,6 +77,7 @@ def test_construct_invalid_field__type_error(key, value): ("machine_id", -1), ("start_time", "not-a-datetime"), ("stop_time", "not-a-datetime"), + ("cc_server", []), ], ) def test_construct_invalid_field__value_error(key, value): @@ -126,7 +127,7 @@ def test_cc_server_set_validated(): a = Agent(**AGENT_SIMPLE_DICT) with pytest.raises(ValueError): - a.cc_server = None + a.cc_server = [] def test_log_contents_set_validated(): diff --git a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py index d03761628..de51fefd2 100644 --- a/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py +++ b/monkey/tests/unit_tests/monkey_island/cc/models/test_agent_registration_message.py @@ -52,7 +52,7 @@ def test_from_serialized(): ("machine_hardware_id", "not-an-int"), ("start_time", None), ("parent_id", 2.1), - ("cc_server", []), + ("cc_server", [1]), ("network_interfaces", "not-a-list"), ], ) @@ -70,6 +70,7 @@ def test_construct_invalid_field__type_error(key, value): ("machine_hardware_id", -1), ("start_time", "not-a-date-time"), ("network_interfaces", [1, "stuff", 3]), + ("cc_server", []), ], ) def test_construct_invalid_field__value_error(key, value): From 1f80eac4b68c929041dbd44726105a9ea4a69a7c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 26 Sep 2022 19:06:35 +0530 Subject: [PATCH 09/12] island: Remove unused import in cc/island_event_handlers/handle_agent_registration.py --- .../cc/island_event_handlers/handle_agent_registration.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2994f03c3..6f3a65669 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,5 +1,5 @@ from contextlib import suppress -from ipaddress import IPv4Address, IPv4Interface +from ipaddress import IPv4Interface from typing import List, Optional from common import AgentRegistrationData From 441c14f15dd9f4d824aa73e420b02dad1a829593 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 27 Sep 2022 12:56:10 +0530 Subject: [PATCH 10/12] Use SocketAddress.from_string() in network/relay/utils.py --- monkey/infection_monkey/network/relay/utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 737a07c9a..0f20a91e3 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -4,7 +4,6 @@ from contextlib import suppress from typing import Dict, Iterable, Iterator, Optional from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT -from common.network.network_utils import address_to_ip_port from common.types import SocketAddress from infection_monkey.island_api_client import ( AbstractIslandAPIClientFactory, @@ -89,8 +88,7 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): def _send_remove_from_waitlist_control_message_to_relay(server: str): - server_address = SocketAddress(*address_to_ip_port(server)) - notify_disconnect(server_address) + notify_disconnect(SocketAddress.from_string(server)) def notify_disconnect(server_address: SocketAddress): From d1199fdab2478e74531b5ba06fefaba139619270 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 27 Sep 2022 13:06:29 +0530 Subject: [PATCH 11/12] Agent: Simplify relay disconnect logic in network/relay/utils.py --- monkey/infection_monkey/network/relay/utils.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 0f20a91e3..86dff8ed6 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -79,18 +79,15 @@ def _check_if_island_server( def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): for i, server in enumerate(servers, start=1): + server_address = SocketAddress.from_string(server) t = create_daemon_thread( - target=_send_remove_from_waitlist_control_message_to_relay, + target=notify_disconnect, name=f"SendRemoveFromWaitlistControlMessageToRelaysThread-{i:02d}", - args=(server,), + args=(server_address,), ) t.start() -def _send_remove_from_waitlist_control_message_to_relay(server: str): - notify_disconnect(SocketAddress.from_string(server)) - - def notify_disconnect(server_address: SocketAddress): """ Tell upstream relay that we no longer need the relay From 9b4de6bab8303a93daef49b8d72abc6361689149 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 27 Sep 2022 13:11:13 +0530 Subject: [PATCH 12/12] UT: Use SocketAddress.from_string() in test_handle_agent_registration.py --- .../island_event_handlers/test_handle_agent_registration.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 61337b022..dff59f5d3 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 @@ -27,12 +27,14 @@ MACHINE = Machine( network_interfaces=[IPv4Interface("192.168.2.2/24")], ) +IP = "192.168.1.1:5000" + AGENT_REGISTRATION_DATA = AgentRegistrationData( id=AGENT_ID, machine_hardware_id=MACHINE.hardware_id, start_time=0, parent_id=None, - cc_server=SocketAddress(ip="192.168.1.1", port="5000"), + cc_server=SocketAddress.from_string(IP), network_interfaces=[IPv4Interface("192.168.1.2/24")], )