From 11da608fe657caacde48c7b2c10e3032c397b876 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 19 Sep 2022 12:09:24 +0200 Subject: [PATCH 01/15] Agent: Modify find_server to use IslandApiClient --- .../infection_monkey/network/relay/utils.py | 21 +++++++++---------- .../network/relay/test_utils.py | 8 +++---- 2 files changed, 14 insertions(+), 15 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index a19f316db..2d2bcc3d7 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -4,11 +4,14 @@ from contextlib import suppress from ipaddress import IPv4Address from typing import Dict, Iterable, Iterator, MutableMapping, Optional -import requests - -from common.common_consts.timeouts import MEDIUM_REQUEST_TIMEOUT from common.network.network_utils import address_to_ip_port from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST +from infection_monkey.transport import IslandApiClient +from infection_monkey.transport.island_api_client_errors import ( + IslandAPIConnectionError, + IslandAPIError, + IslandAPITimeoutError, +) from infection_monkey.utils.threading import ( ThreadSafeIterator, create_daemon_thread, @@ -51,18 +54,14 @@ def _check_if_island_server(server: str) -> bool: logger.debug(f"Trying to connect to server: {server}") try: - requests.get( # noqa: DUO123 - f"https://{server}/api?action=is-up", - verify=False, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) + _ = IslandApiClient(server) return True - except requests.exceptions.ConnectionError as err: + except IslandAPIConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") - except TimeoutError as err: + except IslandAPITimeoutError as err: logger.error(f"Timed out while connecting to server/relay {server}: {err}") - except Exception as err: + except IslandAPIError as err: logger.error( f"Exception encountered when trying to connect to server/relay {server}: {err}" ) 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 ac7eb1b16..58f2f1be7 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,8 +1,8 @@ import pytest -import requests import requests_mock from infection_monkey.network.relay.utils import find_server +from infection_monkey.transport.island_api_client_errors import IslandAPIConnectionError SERVER_1 = "1.1.1.1:12312" SERVER_2 = "2.2.2.2:4321" @@ -16,10 +16,10 @@ servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( "expected_server,server_response_pairs", [ - (None, [(server, {"exc": requests.exceptions.ConnectionError}) for server in servers]), + (None, [(server, {"exc": IslandAPIConnectionError}) for server in servers]), ( SERVER_2, - [(SERVER_1, {"exc": requests.exceptions.ConnectionError})] + [(SERVER_1, {"exc": IslandAPIConnectionError})] + [(server, {"text": ""}) for server in servers[1:]], # type: ignore[dict-item] ), ], @@ -34,7 +34,7 @@ def test_find_server(expected_server, server_response_pairs): def test_find_server__multiple_successes(): with requests_mock.Mocker() as mock: - mock.get(f"https://{SERVER_1}/api?action=is-up", exc=requests.exceptions.ConnectionError) + mock.get(f"https://{SERVER_1}/api?action=is-up", exc=IslandAPIConnectionError) mock.get(f"https://{SERVER_2}/api?action=is-up", text="") mock.get(f"https://{SERVER_3}/api?action=is-up", text="") mock.get(f"https://{SERVER_4}/api?action=is-up", text="") From 8a4666fba22444abdebc8de34495463833531f7c Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 17:24:29 +0530 Subject: [PATCH 02/15] Agent: Fix Island API client import in network/relay/utils.py --- monkey/infection_monkey/network/relay/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 2d2bcc3d7..875d49649 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -6,8 +6,8 @@ from typing import Dict, Iterable, Iterator, MutableMapping, Optional from common.network.network_utils import address_to_ip_port from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST -from infection_monkey.transport import IslandApiClient -from infection_monkey.transport.island_api_client_errors import ( +from infection_monkey.transport import ( + IslandAPIClient, IslandAPIConnectionError, IslandAPIError, IslandAPITimeoutError, @@ -54,7 +54,7 @@ def _check_if_island_server(server: str) -> bool: logger.debug(f"Trying to connect to server: {server}") try: - _ = IslandApiClient(server) + _ = IslandAPIClient(server) return True except IslandAPIConnectionError as err: From 69b26287b6b0c4eabe3ea634de1a777919e7e31b Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 17:24:52 +0530 Subject: [PATCH 03/15] UT: Fix Island API errors' import in network/relay/test_utils.py --- .../unit_tests/infection_monkey/network/relay/test_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 58f2f1be7..ad22e0a4c 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 @@ -2,7 +2,7 @@ import pytest import requests_mock from infection_monkey.network.relay.utils import find_server -from infection_monkey.transport.island_api_client_errors import IslandAPIConnectionError +from infection_monkey.transport import IslandAPIConnectionError SERVER_1 = "1.1.1.1:12312" SERVER_2 = "2.2.2.2:4321" From 9456a30bd923e6c0e095ea585c32a8236e3e11ed Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Mon, 19 Sep 2022 17:26:00 +0530 Subject: [PATCH 04/15] Agent: Remove unnecessary logic in _check_if_island_server() --- 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 875d49649..9a60eb7d5 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -54,7 +54,7 @@ def _check_if_island_server(server: str) -> bool: logger.debug(f"Trying to connect to server: {server}") try: - _ = IslandAPIClient(server) + IslandAPIClient(server) return True except IslandAPIConnectionError as err: From 787af6ae1b83a435586b1afc5195b070bd0e0467 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 08:57:24 -0400 Subject: [PATCH 05/15] Agent: Fix send relay disconnect to unneeded relays Since `find_server()` is parallelized, the iterator was completely exhausted when `send_remove_from_waitlist_control_message_to_relays()` was called, making it effectively a NOOP. --- monkey/infection_monkey/monkey.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 9bb9ef62c..b7eed8b0a 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -138,8 +138,7 @@ class InfectionMonkey: def _get_server(self): logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}") - servers_iterator = (s for s in self._opts.servers) - server = find_server(servers_iterator) + server = find_server(self._opts.servers) if server: logger.info(f"Successfully connected to the island via {server}") else: @@ -147,10 +146,11 @@ class InfectionMonkey: f"Failed to connect to the island via any known servers: {self._opts.servers}" ) - # Note: Since we pass the address for each of our interfaces to the exploited + # 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? - send_remove_from_waitlist_control_message_to_relays(servers_iterator) + servers_to_close = (s for s in self._opts.servers if s != server) + send_remove_from_waitlist_control_message_to_relays(servers_to_close) return server From f4b47f8238310124c555a988e64ad027083a31c4 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 19 Sep 2022 16:44:37 +0200 Subject: [PATCH 06/15] Agent: Use HTTPIslandAPIClient in find_server --- monkey/infection_monkey/network/relay/utils.py | 8 ++++---- .../infection_monkey/network/relay/test_utils.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 9a60eb7d5..c2ffc64ba 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -5,13 +5,13 @@ from ipaddress import IPv4Address from typing import Dict, Iterable, Iterator, MutableMapping, Optional from common.network.network_utils import address_to_ip_port -from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST -from infection_monkey.transport import ( - IslandAPIClient, +from infection_monkey.island_api_client import ( + HTTPIslandAPIClient, IslandAPIConnectionError, IslandAPIError, IslandAPITimeoutError, ) +from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST from infection_monkey.utils.threading import ( ThreadSafeIterator, create_daemon_thread, @@ -54,7 +54,7 @@ def _check_if_island_server(server: str) -> bool: logger.debug(f"Trying to connect to server: {server}") try: - IslandAPIClient(server) + HTTPIslandAPIClient(server) return True except IslandAPIConnectionError as err: 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 ad22e0a4c..c4e69bd1a 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,8 +1,8 @@ import pytest import requests_mock +from infection_monkey.island_api_client import IslandAPIConnectionError from infection_monkey.network.relay.utils import find_server -from infection_monkey.transport import IslandAPIConnectionError SERVER_1 = "1.1.1.1:12312" SERVER_2 = "2.2.2.2:4321" From bc19b5ea93194a2b321a55b1518ca13f4355c7b7 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Mon, 19 Sep 2022 17:24:00 +0200 Subject: [PATCH 07/15] Agent: Modify find_server to return tuple of server and IIslandAPIClient --- monkey/infection_monkey/monkey.py | 11 +++++---- .../infection_monkey/network/relay/utils.py | 24 +++++++++++-------- .../network/relay/test_utils.py | 13 ++++++---- 3 files changed, 29 insertions(+), 19 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b7eed8b0a..b1a222af0 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -5,7 +5,7 @@ import subprocess import sys from ipaddress import IPv4Address, IPv4Interface from pathlib import Path, WindowsPath -from typing import List +from typing import List, Tuple from pubsub.core import Publisher @@ -45,6 +45,7 @@ from infection_monkey.exploit.sshexec import SSHExploiter from infection_monkey.exploit.wmiexec import WmiExploiter from infection_monkey.exploit.zerologon import ZerologonExploiter from infection_monkey.i_puppet import IPuppet, PluginType +from infection_monkey.island_api_client import IIslandAPIClient from infection_monkey.master import AutomatedMaster from infection_monkey.master.control_channel import ControlChannel from infection_monkey.model import VictimHostFactory @@ -110,7 +111,7 @@ class InfectionMonkey: self._opts = self._get_arguments(args) # TODO: Revisit variable names - server = self._get_server() + server, island_api_client = self._get_server() # 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) @@ -136,9 +137,9 @@ class InfectionMonkey: return opts - def _get_server(self): + def _get_server(self) -> Tuple[str, IIslandAPIClient]: logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}") - server = find_server(self._opts.servers) + server, island_api_client = find_server(self._opts.servers) if server: logger.info(f"Successfully connected to the island via {server}") else: @@ -152,7 +153,7 @@ class InfectionMonkey: servers_to_close = (s for s in self._opts.servers if s != server) send_remove_from_waitlist_control_message_to_relays(servers_to_close) - return server + return server, island_api_client @staticmethod def _log_arguments(args): diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index c2ffc64ba..acbd9cc46 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -2,11 +2,12 @@ import logging import socket from contextlib import suppress from ipaddress import IPv4Address -from typing import Dict, Iterable, Iterator, MutableMapping, Optional +from typing import Dict, Iterable, Iterator, MutableMapping, Optional, Tuple from common.network.network_utils import address_to_ip_port from infection_monkey.island_api_client import ( HTTPIslandAPIClient, + IIslandAPIClient, IslandAPIConnectionError, IslandAPIError, IslandAPITimeoutError, @@ -25,10 +26,10 @@ logger = logging.getLogger(__name__) NUM_FIND_SERVER_WORKERS = 32 -def find_server(servers: Iterable[str]) -> Optional[str]: +def find_server(servers: Iterable[str]) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: server_list = list(servers) server_iterator = ThreadSafeIterator(server_list.__iter__()) - server_results: Dict[str, bool] = {} + server_results: Dict[str, Tuple[bool, IIslandAPIClient]] = {} run_worker_threads( _find_island_server, @@ -39,24 +40,27 @@ def find_server(servers: Iterable[str]) -> Optional[str]: for server in server_list: if server_results[server]: - return server + island_api_client = server_results[server] + return server, island_api_client - return None + return (None, None) -def _find_island_server(servers: Iterator[str], server_status: MutableMapping[str, bool]): +def _find_island_server( + servers: Iterator[str], server_status: MutableMapping[str, Optional[IIslandAPIClient]] +): with suppress(StopIteration): server = next(servers) server_status[server] = _check_if_island_server(server) -def _check_if_island_server(server: str) -> bool: +def _check_if_island_server(server: str) -> IIslandAPIClient: logger.debug(f"Trying to connect to server: {server}") try: - HTTPIslandAPIClient(server) + island_api_client = HTTPIslandAPIClient(server) - return True + return island_api_client except IslandAPIConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") except IslandAPITimeoutError as err: @@ -66,7 +70,7 @@ def _check_if_island_server(server: str) -> bool: f"Exception encountered when trying to connect to server/relay {server}: {err}" ) - return False + return None def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): 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 c4e69bd1a..4979cce43 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,7 +1,7 @@ import pytest import requests_mock -from infection_monkey.island_api_client import IslandAPIConnectionError +from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIConnectionError from infection_monkey.network.relay.utils import find_server SERVER_1 = "1.1.1.1:12312" @@ -14,7 +14,7 @@ servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( - "expected_server,server_response_pairs", + "expected_server, server_response_pairs", [ (None, [(server, {"exc": IslandAPIConnectionError}) for server in servers]), ( @@ -29,7 +29,9 @@ def test_find_server(expected_server, server_response_pairs): for server, response in server_response_pairs: mock.get(f"https://{server}/api?action=is-up", **response) - assert find_server(servers) is expected_server + actual_server, _ = find_server(servers) + + assert actual_server is expected_server def test_find_server__multiple_successes(): @@ -39,4 +41,7 @@ def test_find_server__multiple_successes(): mock.get(f"https://{SERVER_3}/api?action=is-up", text="") mock.get(f"https://{SERVER_4}/api?action=is-up", text="") - assert find_server(servers) == SERVER_2 + actual_server, actual_island_api_client = find_server(servers) + + assert actual_server == SERVER_2 + assert isinstance(actual_island_api_client, IIslandAPIClient) From db75806a083478bed864fdf5f9a43b218b7f6b7f Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 12:13:03 -0400 Subject: [PATCH 08/15] Agent: Rename _get_server() -> _connect_to_island_api() --- 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 b1a222af0..034f47295 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -111,7 +111,7 @@ class InfectionMonkey: self._opts = self._get_arguments(args) # TODO: Revisit variable names - server, island_api_client = self._get_server() + server, 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) @@ -137,7 +137,7 @@ class InfectionMonkey: return opts - def _get_server(self) -> Tuple[str, IIslandAPIClient]: + def _connect_to_island_api(self) -> Tuple[str, IIslandAPIClient]: logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}") server, island_api_client = find_server(self._opts.servers) if server: From 6563be82226801eba37226024056d8d9d771979a Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 12:16:20 -0400 Subject: [PATCH 09/15] Agent: Remove unnecessary local variable --- 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 acbd9cc46..c964ad8c0 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -58,9 +58,7 @@ def _check_if_island_server(server: str) -> IIslandAPIClient: logger.debug(f"Trying to connect to server: {server}") try: - island_api_client = HTTPIslandAPIClient(server) - - return island_api_client + return HTTPIslandAPIClient(server) except IslandAPIConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") except IslandAPITimeoutError as err: From c6a5e294dff9f68d1fc2fece515975baea20fecd Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 13:57:25 -0400 Subject: [PATCH 10/15] Agent: Add a timeout to notify disconnect socket --- monkey/infection_monkey/network/relay/utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index c964ad8c0..6d4944e3b 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -4,6 +4,7 @@ from contextlib import suppress from ipaddress import IPv4Address from typing import Dict, Iterable, Iterator, MutableMapping, Optional, Tuple +from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT from common.network.network_utils import address_to_ip_port from infection_monkey.island_api_client import ( HTTPIslandAPIClient, @@ -94,6 +95,8 @@ def notify_disconnect(server_ip: IPv4Address, server_port: int): :param server_port: The port of the server to notify. """ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: + d_socket.settimeout(LONG_REQUEST_TIMEOUT) + try: d_socket.connect((server_ip, server_port)) d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST) From 51ecf52d4bfc4f347efe8e53ee3491d2ce57c9b9 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 14:01:08 -0400 Subject: [PATCH 11/15] Agent: Add an ID (int) to relay control message thread name --- monkey/infection_monkey/network/relay/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 6d4944e3b..cdb44ced0 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -73,10 +73,10 @@ def _check_if_island_server(server: str) -> IIslandAPIClient: def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): - for server in servers: + for i, server in enumerate(servers, start=1): t = create_daemon_thread( target=_send_remove_from_waitlist_control_message_to_relay, - name="SendRemoveFromWaitlistControlMessageToRelaysThread", + name=f"SendRemoveFromWaitlistControlMessageToRelaysThread-{i:02d}", args=(server,), ) t.start() From 9ea291a7fa4462643e3b87356cd165d83a412f0c Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 14:01:51 -0400 Subject: [PATCH 12/15] Agent: Fix crash when connecting to IPv4Address socket.connect() needs a string, not IPv4Address, otherwise the thread will crash. --- 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 cdb44ced0..4a704bf08 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -98,7 +98,7 @@ def notify_disconnect(server_ip: IPv4Address, server_port: int): d_socket.settimeout(LONG_REQUEST_TIMEOUT) try: - d_socket.connect((server_ip, server_port)) + d_socket.connect((str(server_ip), server_port)) d_socket.sendall(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST) logger.info(f"Control message was sent to the server/relay {server_ip}:{server_port}") except OSError as err: From 2ebb7621e37230fde7698d74005532b6c4ad0b1e Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 14:05:34 -0400 Subject: [PATCH 13/15] Agent: Fix server selection logic --- monkey/infection_monkey/monkey.py | 19 ++++++++++++++++--- .../infection_monkey/network/relay/utils.py | 11 +++-------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 034f47295..33042d2cf 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -5,7 +5,7 @@ import subprocess import sys from ipaddress import IPv4Address, IPv4Interface from pathlib import Path, WindowsPath -from typing import List, Tuple +from typing import List, Mapping, Optional, Tuple from pubsub.core import Publisher @@ -137,9 +137,13 @@ 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[str, IIslandAPIClient]: logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}") - server, island_api_client = find_server(self._opts.servers) + server_clients = find_server(self._opts.servers) + + server, island_api_client = self._select_server(server_clients) + if server: logger.info(f"Successfully connected to the island via {server}") else: @@ -150,11 +154,20 @@ class InfectionMonkey: # 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._opts.servers if s != server) + 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 + def _select_server( + self, server_clients: Mapping[str, IIslandAPIClient] + ) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: + for server in self._opts.servers: + if server_clients[server]: + return server, server_clients[server] + + return None, None + @staticmethod def _log_arguments(args): arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()]) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 4a704bf08..e31c9ac97 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -2,7 +2,7 @@ import logging import socket from contextlib import suppress from ipaddress import IPv4Address -from typing import Dict, Iterable, Iterator, MutableMapping, Optional, Tuple +from typing import Dict, Iterable, Iterator, Mapping, MutableMapping, Optional, Tuple from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT from common.network.network_utils import address_to_ip_port @@ -27,7 +27,7 @@ logger = logging.getLogger(__name__) NUM_FIND_SERVER_WORKERS = 32 -def find_server(servers: Iterable[str]) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: +def find_server(servers: Iterable[str]) -> Mapping[str, Optional[IIslandAPIClient]]: server_list = list(servers) server_iterator = ThreadSafeIterator(server_list.__iter__()) server_results: Dict[str, Tuple[bool, IIslandAPIClient]] = {} @@ -39,12 +39,7 @@ def find_server(servers: Iterable[str]) -> Tuple[Optional[str], Optional[IIsland num_workers=NUM_FIND_SERVER_WORKERS, ) - for server in server_list: - if server_results[server]: - island_api_client = server_results[server] - return server, island_api_client - - return (None, None) + return server_results def _find_island_server( From 753ac739b015ba9abba7b81ed224aaa727379ce2 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 14:20:40 -0400 Subject: [PATCH 14/15] Agent: rename find_servers() -. find_available_island_apis() --- monkey/infection_monkey/monkey.py | 4 ++-- monkey/infection_monkey/network/relay/utils.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 33042d2cf..680db9f6f 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -53,7 +53,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 ( - find_server, + find_available_island_apis, notify_disconnect, send_remove_from_waitlist_control_message_to_relays, ) @@ -140,7 +140,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[str, IIslandAPIClient]: logger.debug(f"Trying to wake up with servers: {', '.join(self._opts.servers)}") - server_clients = find_server(self._opts.servers) + server_clients = find_available_island_apis(self._opts.servers) server, island_api_client = self._select_server(server_clients) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index e31c9ac97..451fd65e7 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -27,7 +27,7 @@ logger = logging.getLogger(__name__) NUM_FIND_SERVER_WORKERS = 32 -def find_server(servers: Iterable[str]) -> Mapping[str, Optional[IIslandAPIClient]]: +def find_available_island_apis(servers: Iterable[str]) -> Mapping[str, Optional[IIslandAPIClient]]: server_list = list(servers) server_iterator = ThreadSafeIterator(server_list.__iter__()) server_results: Dict[str, Tuple[bool, IIslandAPIClient]] = {} From b9576db426ef799d6a8782e2569dad5c77942cb8 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 19 Sep 2022 14:21:17 -0400 Subject: [PATCH 15/15] UT: Fix broken tests for find_available_island_apis() --- .../network/relay/test_utils.py | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) 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 4979cce43..dd00ebc94 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 @@ -2,7 +2,7 @@ import pytest import requests_mock from infection_monkey.island_api_client import IIslandAPIClient, IslandAPIConnectionError -from infection_monkey.network.relay.utils import find_server +from infection_monkey.network.relay.utils import find_available_island_apis SERVER_1 = "1.1.1.1:12312" SERVER_2 = "2.2.2.2:4321" @@ -14,34 +14,42 @@ servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( - "expected_server, server_response_pairs", + "expected_available_servers, server_response_pairs", [ - (None, [(server, {"exc": IslandAPIConnectionError}) for server in servers]), + ([], [(server, {"exc": IslandAPIConnectionError}) for server in servers]), ( - SERVER_2, + servers[1:], [(SERVER_1, {"exc": IslandAPIConnectionError})] + [(server, {"text": ""}) for server in servers[1:]], # type: ignore[dict-item] ), ], ) -def test_find_server(expected_server, server_response_pairs): +def test_find_available_island_apis(expected_available_servers, server_response_pairs): with requests_mock.Mocker() as mock: for server, response in server_response_pairs: mock.get(f"https://{server}/api?action=is-up", **response) - actual_server, _ = find_server(servers) + available_apis = find_available_island_apis(servers) - assert actual_server is expected_server + 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 + else: + assert island_api_client is None -def test_find_server__multiple_successes(): +def test_find_available_island_apis__multiple_successes(): + available_servers = [SERVER_2, SERVER_3] with requests_mock.Mocker() as mock: mock.get(f"https://{SERVER_1}/api?action=is-up", exc=IslandAPIConnectionError) - mock.get(f"https://{SERVER_2}/api?action=is-up", text="") - mock.get(f"https://{SERVER_3}/api?action=is-up", text="") - mock.get(f"https://{SERVER_4}/api?action=is-up", text="") + for server in available_servers: + mock.get(f"https://{server}/api?action=is-up", text="") - actual_server, actual_island_api_client = find_server(servers) + available_apis = find_available_island_apis(servers) - assert actual_server == SERVER_2 - assert isinstance(actual_island_api_client, IIslandAPIClient) + 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)