From 110542eeb851f653ba92f6d66ca481867996051e Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 13:27:08 +0000 Subject: [PATCH] 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)