From c4804f06a9d63bcb0ae1450350d6fd374a5ae13d Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 23 Sep 2022 21:28:39 +0000 Subject: [PATCH 1/7] Agent: Use SocketAddress in IIslandAPIClient --- .../http_island_api_client.py | 6 +- .../island_api_client/i_island_api_client.py | 3 +- monkey/infection_monkey/monkey.py | 15 ++--- .../infection_monkey/network/relay/utils.py | 34 +++++++---- .../test_http_island_api_client.py | 2 +- .../network/relay/test_utils.py | 59 +++++++++++++++---- 6 files changed, 82 insertions(+), 37 deletions(-) diff --git a/monkey/infection_monkey/island_api_client/http_island_api_client.py b/monkey/infection_monkey/island_api_client/http_island_api_client.py index 09fcbf762..624a2c504 100644 --- a/monkey/infection_monkey/island_api_client/http_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/http_island_api_client.py @@ -16,6 +16,7 @@ from common.common_consts.timeouts import ( SHORT_REQUEST_TIMEOUT, ) from common.credentials import Credentials +from common.types import SocketAddress from . import ( AbstractIslandAPIClientFactory, @@ -79,7 +80,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): @handle_island_errors def connect( self, - island_server: str, + island_server: SocketAddress, ): response = requests.get( # noqa: DUO123 f"https://{island_server}/api?action=is-up", @@ -88,8 +89,7 @@ class HTTPIslandAPIClient(IIslandAPIClient): ) response.raise_for_status() - self._island_server = island_server - self._api_url = f"https://{self._island_server}/api" + self._api_url = f"https://{island_server}/api" @handle_island_errors def send_log(self, log_contents: str): diff --git a/monkey/infection_monkey/island_api_client/i_island_api_client.py b/monkey/infection_monkey/island_api_client/i_island_api_client.py index e29c19115..3fd83daaa 100644 --- a/monkey/infection_monkey/island_api_client/i_island_api_client.py +++ b/monkey/infection_monkey/island_api_client/i_island_api_client.py @@ -5,6 +5,7 @@ from common import AgentRegistrationData, AgentSignals, OperatingSystem from common.agent_configuration import AgentConfiguration from common.agent_events import AbstractAgentEvent from common.credentials import Credentials +from common.types import SocketAddress class IIslandAPIClient(ABC): @@ -13,7 +14,7 @@ class IIslandAPIClient(ABC): """ @abstractmethod - def connect(self, island_server: str): + def connect(self, island_server: SocketAddress): """ Connect to the island's API diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 80f6d9735..923f73da2 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -5,7 +5,7 @@ import subprocess import sys from ipaddress import IPv4Interface from pathlib import Path, WindowsPath -from typing import List, Mapping, Optional, Tuple +from typing import List, Optional, Tuple from pubsub.core import Publisher @@ -55,6 +55,7 @@ from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_free_tcp_port from infection_monkey.network.relay import TCPRelay from infection_monkey.network.relay.utils import ( + IslandAPISearchResults, find_available_island_apis, notify_disconnect, send_remove_from_waitlist_control_message_to_relays, @@ -156,7 +157,7 @@ class InfectionMonkey: def _connect_to_island_api(self) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: logger.debug(f"Trying to wake up with servers: {', '.join(self._server_strings)}") server_clients = find_available_island_apis( - self._server_strings, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) + self._opts.servers, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) ) server, island_api_client = self._select_server(server_clients) @@ -189,11 +190,11 @@ class InfectionMonkey: self._island_api_client.register_agent(agent_registration_data) def _select_server( - self, server_clients: Mapping[str, Optional[IIslandAPIClient]] - ) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: - for server in self._server_strings: - if server_clients[server]: - return server, server_clients[server] + self, server_clients: IslandAPISearchResults + ) -> Tuple[Optional[SocketAddress], Optional[IIslandAPIClient]]: + for result in server_clients: + if result.client is not None: + return result.server, result.client return None, None diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 86dff8ed6..9dbb9cf64 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,7 +1,8 @@ import logging import socket from contextlib import suppress -from typing import Dict, Iterable, Iterator, Optional +from dataclasses import dataclass +from typing import Dict, Iterable, Iterator, List, Optional, Tuple from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT from common.types import SocketAddress @@ -25,38 +26,47 @@ logger = logging.getLogger(__name__) # practical purposes. Revisit this if it's not. NUM_FIND_SERVER_WORKERS = 32 -IslandAPISearchResults = Dict[str, Optional[IIslandAPIClient]] + +@dataclass +class IslandAPISearchResult: + server: SocketAddress + client: Optional[IIslandAPIClient] + + +IslandAPISearchResults = List[IslandAPISearchResult] def find_available_island_apis( - servers: Iterable[str], island_api_client_factory: AbstractIslandAPIClientFactory + servers: Iterable[SocketAddress], island_api_client_factory: AbstractIslandAPIClientFactory ) -> IslandAPISearchResults: server_list = list(servers) - server_iterator = ThreadSafeIterator(server_list.__iter__()) - server_results: IslandAPISearchResults = {} + server_iterator = ThreadSafeIterator(enumerate(server_list.__iter__())) + results: Dict[int, IslandAPISearchResult] = {} run_worker_threads( _find_island_server, "FindIslandServer", - args=(server_iterator, server_results, island_api_client_factory), + args=(server_iterator, results, island_api_client_factory), num_workers=NUM_FIND_SERVER_WORKERS, ) - return server_results + return [results[i] for i in sorted(results.keys())] def _find_island_server( - servers: Iterator[str], - server_results: IslandAPISearchResults, + servers: Iterator[Tuple[int, SocketAddress]], + server_results: Dict[int, IslandAPISearchResult], island_api_client_factory: AbstractIslandAPIClientFactory, ): with suppress(StopIteration): - server = next(servers) - server_results[server] = _check_if_island_server(server, island_api_client_factory) + index, server = next(servers) + server_results[index] = IslandAPISearchResult( + server, _check_if_island_server(server, island_api_client_factory) + ) def _check_if_island_server( - server: str, island_api_client_factory: AbstractIslandAPIClientFactory + server: SocketAddress, island_api_client_factory: AbstractIslandAPIClientFactory ) -> Optional[IIslandAPIClient]: logger.debug(f"Trying to connect to server: {server}") 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 56a480082..c10c405d8 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 @@ -30,7 +30,7 @@ AGENT_REGISTRATION = AgentRegistrationData( machine_hardware_id=1, start_time=0, parent_id=None, - cc_server=SERVER, + cc_server=str(SERVER), network_interfaces=[], ) diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py index affb61aff..4ace390c0 100644 --- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py @@ -1,18 +1,21 @@ +from typing import Callable, Optional + import pytest import requests_mock from common.agent_event_serializers import AgentEventSerializerRegistry +from common.types import SocketAddress from infection_monkey.island_api_client import ( HTTPIslandAPIClientFactory, IIslandAPIClient, IslandAPIConnectionError, ) -from infection_monkey.network.relay.utils import find_available_island_apis +from infection_monkey.network.relay.utils import IslandAPISearchResult, find_available_island_apis -SERVER_1 = "1.1.1.1:12312" -SERVER_2 = "2.2.2.2:4321" -SERVER_3 = "3.3.3.3:3142" -SERVER_4 = "4.4.4.4:5000" +SERVER_1 = SocketAddress(ip="1.1.1.1", port=12312) +SERVER_2 = SocketAddress(ip="2.2.2.2", port=4321) +SERVER_3 = SocketAddress(ip="3.3.3.3", port=3142) +SERVER_4 = SocketAddress(ip="4.4.4.4", port=5000) servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @@ -45,11 +48,41 @@ def test_find_available_island_apis( assert len(available_apis) == len(server_response_pairs) - for server, island_api_client in available_apis.items(): - if server in expected_available_servers: - assert island_api_client is not None + for result in available_apis: + if result.server in expected_available_servers: + assert result.client is not None else: - assert island_api_client is None + assert result.client is None + + +def test_find_available_island_apis__preserves_input_order(island_api_client_factory): + available_servers = [SERVER_2, SERVER_3] + + with requests_mock.Mocker() as mock: + mock.get(f"https://{SERVER_1}/api?action=is-up", exc=IslandAPIConnectionError) + for server in available_servers: + mock.get(f"https://{server}/api?action=is-up", text="") + available_apis = find_available_island_apis(servers, island_api_client_factory) + + for index in range(len(servers)): + assert available_apis[index].server == servers[index] + + +def _is_none(value) -> bool: + return value is None + + +def _is_island_client(value) -> bool: + return isinstance(value, IIslandAPIClient) + + +def _assert_server_and_predicate( + result: IslandAPISearchResult, + server: SocketAddress, + predicate: Callable[[Optional[IIslandAPIClient]], bool], +): + assert result.server == server + assert predicate(result.client) def test_find_available_island_apis__multiple_successes(island_api_client_factory): @@ -61,7 +94,7 @@ def test_find_available_island_apis__multiple_successes(island_api_client_factor available_apis = find_available_island_apis(servers, island_api_client_factory) - assert available_apis[SERVER_1] is None - assert available_apis[SERVER_4] is None - for server in available_servers: - assert isinstance(available_apis[server], IIslandAPIClient) + _assert_server_and_predicate(available_apis[0], SERVER_1, _is_none) + _assert_server_and_predicate(available_apis[1], SERVER_2, _is_island_client) + _assert_server_and_predicate(available_apis[2], SERVER_3, _is_island_client) + _assert_server_and_predicate(available_apis[3], SERVER_4, _is_none) From 110542eeb851f653ba92f6d66ca481867996051e Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 13:27:08 +0000 Subject: [PATCH 2/7] Common: Make SocketAddress hashable --- monkey/common/types.py | 3 ++ monkey/infection_monkey/monkey.py | 17 ++++--- .../infection_monkey/network/relay/utils.py | 29 ++++------- .../network/relay/test_utils.py | 50 ++++--------------- 4 files changed, 31 insertions(+), 68 deletions(-) diff --git a/monkey/common/types.py b/monkey/common/types.py index 0c14aeec9..c1263a466 100644 --- a/monkey/common/types.py +++ b/monkey/common/types.py @@ -41,5 +41,8 @@ class SocketAddress(InfectionMonkeyBaseModel): raise ValueError("SocketAddress requires a port") return SocketAddress(ip=IPv4Address(ip), port=int(port)) + def __hash__(self): + return hash(str(self)) + def __str__(self): return f"{self.ip}:{self.port}" diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 923f73da2..d84693ec7 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -115,7 +115,6 @@ class InfectionMonkey: self._singleton = SystemSingleton() self._opts = self._get_arguments(args) - self._server_strings = [str(s) for s in self._opts.servers] self._agent_event_serializer_registry = self._setup_agent_event_serializers() @@ -155,7 +154,7 @@ class InfectionMonkey: # TODO: By the time we finish 2292, _connect_to_island_api() may not need to return `server` def _connect_to_island_api(self) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: - logger.debug(f"Trying to wake up with servers: {', '.join(self._server_strings)}") + logger.debug(f"Trying to wake up with servers: {', '.join(map(str, self._opts.servers))}") server_clients = find_available_island_apis( self._opts.servers, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) ) @@ -166,13 +165,14 @@ class InfectionMonkey: logger.info(f"Successfully connected to the island via {server}") else: raise Exception( - f"Failed to connect to the island via any known servers: {self._server_strings}" + "Failed to connect to the island via any known servers: " + f"[{', '.join(map(str, self._opts.servers))}]" ) # NOTE: Since we pass the address for each of our interfaces to the exploited # machines, is it possible for a machine to unintentionally unregister itself from the # relay if it is able to connect to the relay over multiple interfaces? - servers_to_close = (s for s in self._server_strings if s != server and server_clients[s]) + servers_to_close = (s for s in self._opts.servers if s != server and server_clients[s]) send_remove_from_waitlist_control_message_to_relays(servers_to_close) return server, island_api_client @@ -192,9 +192,9 @@ class InfectionMonkey: def _select_server( self, server_clients: IslandAPISearchResults ) -> Tuple[Optional[SocketAddress], Optional[IIslandAPIClient]]: - for result in server_clients: - if result.client is not None: - return result.server, result.client + for server in self._opts.servers: + if server_clients[server] is not None: + return server, server_clients[server] return None, None @@ -261,8 +261,9 @@ class InfectionMonkey: return agent_event_serializer_registry def _build_server_list(self, relay_port: int): + my_servers = [str(s) for s in self._opts.servers] relay_servers = [f"{ip}:{relay_port}" for ip in get_my_ip_addresses()] - return self._server_strings + relay_servers + return my_servers + relay_servers def _build_master(self, relay_port: int): servers = self._build_server_list(relay_port) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 9dbb9cf64..247a5c50a 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,8 +1,7 @@ import logging import socket from contextlib import suppress -from dataclasses import dataclass -from typing import Dict, Iterable, Iterator, List, Optional, Tuple +from typing import Dict, Iterable, Iterator, Optional from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT from common.types import SocketAddress @@ -27,42 +26,34 @@ logger = logging.getLogger(__name__) NUM_FIND_SERVER_WORKERS = 32 -@dataclass -class IslandAPISearchResult: - server: SocketAddress - client: Optional[IIslandAPIClient] - - -IslandAPISearchResults = List[IslandAPISearchResult] +IslandAPISearchResults = Dict[SocketAddress, Optional[IIslandAPIClient]] def find_available_island_apis( servers: Iterable[SocketAddress], island_api_client_factory: AbstractIslandAPIClientFactory ) -> IslandAPISearchResults: server_list = list(servers) - server_iterator = ThreadSafeIterator(enumerate(server_list.__iter__())) - results: Dict[int, IslandAPISearchResult] = {} + server_iterator = ThreadSafeIterator(server_list.__iter__()) + server_results: IslandAPISearchResults = {} run_worker_threads( _find_island_server, "FindIslandServer", - args=(server_iterator, results, island_api_client_factory), + args=(server_iterator, server_results, island_api_client_factory), num_workers=NUM_FIND_SERVER_WORKERS, ) - return [results[i] for i in sorted(results.keys())] + return server_results def _find_island_server( - servers: Iterator[Tuple[int, SocketAddress]], - server_results: Dict[int, IslandAPISearchResult], + servers: Iterator[SocketAddress], + server_results: IslandAPISearchResults, island_api_client_factory: AbstractIslandAPIClientFactory, ): with suppress(StopIteration): - index, server = next(servers) - server_results[index] = IslandAPISearchResult( - server, _check_if_island_server(server, island_api_client_factory) - ) + server = next(servers) + server_results[server] = _check_if_island_server(server, island_api_client_factory) def _check_if_island_server( diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py index 4ace390c0..6a7aba04d 100644 --- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py @@ -1,5 +1,3 @@ -from typing import Callable, Optional - import pytest import requests_mock @@ -10,7 +8,7 @@ from infection_monkey.island_api_client import ( IIslandAPIClient, IslandAPIConnectionError, ) -from infection_monkey.network.relay.utils import IslandAPISearchResult, find_available_island_apis +from infection_monkey.network.relay.utils import find_available_island_apis SERVER_1 = SocketAddress(ip="1.1.1.1", port=12312) SERVER_2 = SocketAddress(ip="2.2.2.2", port=4321) @@ -48,41 +46,11 @@ def test_find_available_island_apis( assert len(available_apis) == len(server_response_pairs) - for result in available_apis: - if result.server in expected_available_servers: - assert result.client is not None + for server, island_api_client in available_apis.items(): + if server in expected_available_servers: + assert island_api_client is not None else: - assert result.client is None - - -def test_find_available_island_apis__preserves_input_order(island_api_client_factory): - available_servers = [SERVER_2, SERVER_3] - - with requests_mock.Mocker() as mock: - mock.get(f"https://{SERVER_1}/api?action=is-up", exc=IslandAPIConnectionError) - for server in available_servers: - mock.get(f"https://{server}/api?action=is-up", text="") - available_apis = find_available_island_apis(servers, island_api_client_factory) - - for index in range(len(servers)): - assert available_apis[index].server == servers[index] - - -def _is_none(value) -> bool: - return value is None - - -def _is_island_client(value) -> bool: - return isinstance(value, IIslandAPIClient) - - -def _assert_server_and_predicate( - result: IslandAPISearchResult, - server: SocketAddress, - predicate: Callable[[Optional[IIslandAPIClient]], bool], -): - assert result.server == server - assert predicate(result.client) + assert island_api_client is None def test_find_available_island_apis__multiple_successes(island_api_client_factory): @@ -94,7 +62,7 @@ def test_find_available_island_apis__multiple_successes(island_api_client_factor available_apis = find_available_island_apis(servers, island_api_client_factory) - _assert_server_and_predicate(available_apis[0], SERVER_1, _is_none) - _assert_server_and_predicate(available_apis[1], SERVER_2, _is_island_client) - _assert_server_and_predicate(available_apis[2], SERVER_3, _is_island_client) - _assert_server_and_predicate(available_apis[3], SERVER_4, _is_none) + assert available_apis[SERVER_1] is None + assert available_apis[SERVER_4] is None + for server in available_servers: + assert isinstance(available_apis[server], IIslandAPIClient) From af8d3937be704018398ae50aeccff4763ecadfdc Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 15:28:23 +0000 Subject: [PATCH 3/7] Agent: Use correct return type for _connect_to_island_api --- monkey/infection_monkey/monkey.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index d84693ec7..797abb3cb 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -119,12 +119,13 @@ class InfectionMonkey: self._agent_event_serializer_registry = self._setup_agent_event_serializers() server, self._island_api_client = self._connect_to_island_api() - self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) + self._cmd_island_ip = server.ip + self._cmd_island_port = server.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 + server_address=str(server), island_api_client=self._island_api_client ) self._control_channel = ControlChannel(server, get_agent_id(), self._island_api_client) self._register_agent(self._island_address) @@ -153,7 +154,7 @@ class InfectionMonkey: return opts # TODO: By the time we finish 2292, _connect_to_island_api() may not need to return `server` - def _connect_to_island_api(self) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: + def _connect_to_island_api(self) -> Tuple[Optional[SocketAddress], Optional[IIslandAPIClient]]: logger.debug(f"Trying to wake up with servers: {', '.join(map(str, self._opts.servers))}") server_clients = find_available_island_apis( self._opts.servers, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) @@ -278,7 +279,6 @@ class InfectionMonkey: self._subscribe_events( event_queue, propagation_credentials_repository, - self._control_client.server_address, self._agent_event_serializer_registry, ) @@ -305,7 +305,6 @@ class InfectionMonkey: self, event_queue: IAgentEventQueue, propagation_credentials_repository: IPropagationCredentialsRepository, - server_address: str, agent_event_serializer_registry: AgentEventSerializerRegistry, ): event_queue.subscribe_type( From aeef2cdcbebfad30d69a16115ff7fd0f571828b3 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 16:14:27 +0000 Subject: [PATCH 4/7] Agent: Update send_remove_from_waitlist_control_message_to_relays Update send_remove_from_waitlist_control_message_to_relays to use SocketAddress --- monkey/infection_monkey/network/relay/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 247a5c50a..fa732f7ff 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -78,7 +78,7 @@ def _check_if_island_server( return None -def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): +def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[SocketAddress]): for i, server in enumerate(servers, start=1): server_address = SocketAddress.from_string(server) t = create_daemon_thread( From 6d63f3c378229f7d0e66948cab42702e4825b5b8 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 27 Sep 2022 16:17:04 +0530 Subject: [PATCH 5/7] Agent: Fix some logic caused by rebase --- monkey/infection_monkey/monkey.py | 2 +- monkey/infection_monkey/network/relay/utils.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 797abb3cb..e336e6db9 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -127,7 +127,7 @@ class InfectionMonkey: self._control_client = ControlClient( server_address=str(server), island_api_client=self._island_api_client ) - self._control_channel = ControlChannel(server, get_agent_id(), self._island_api_client) + self._control_channel = ControlChannel(str(server), get_agent_id(), self._island_api_client) self._register_agent(self._island_address) # TODO Refactor the telemetry messengers to accept control client diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index fa732f7ff..ca617306e 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -80,11 +80,10 @@ def _check_if_island_server( def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[SocketAddress]): for i, server in enumerate(servers, start=1): - server_address = SocketAddress.from_string(server) t = create_daemon_thread( target=notify_disconnect, name=f"SendRemoveFromWaitlistControlMessageToRelaysThread-{i:02d}", - args=(server_address,), + args=(server,), ) t.start() From 01f1d62272e4e6996cde9c3cdea4ee99594b9ab4 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 27 Sep 2022 16:19:23 +0530 Subject: [PATCH 6/7] UT: Simplify logic in test data in test_http_island_api_client.py --- .../island_api_client/test_http_island_api_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c10c405d8..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 @@ -30,7 +30,7 @@ AGENT_REGISTRATION = AgentRegistrationData( machine_hardware_id=1, start_time=0, parent_id=None, - cc_server=str(SERVER), + cc_server=SERVER, network_interfaces=[], ) From 4f3a8a5b2fce9d32b92979877a8da329192e1e9f Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Tue, 27 Sep 2022 17:17:09 +0530 Subject: [PATCH 7/7] Agent: Simplify logic in _build_server_list() in monkey/infection_monkey/monkey.py Co-authored-by: VakarisZ <36815064+VakarisZ@users.noreply.github.com> --- monkey/infection_monkey/monkey.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index e336e6db9..d8e4962d9 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -262,7 +262,7 @@ class InfectionMonkey: return agent_event_serializer_registry def _build_server_list(self, relay_port: int): - my_servers = [str(s) for s in self._opts.servers] + my_servers = map(str, self._opts.servers) relay_servers = [f"{ip}:{relay_port}" for ip in get_my_ip_addresses()] return my_servers + relay_servers