From ac058c77883c3b5ab9d94a9b482384129da81a4d Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 13:20:20 +0200 Subject: [PATCH 01/31] Agent: Add TODO to rework address_to_ip_port to use list of servers --- monkey/infection_monkey/monkey.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index f8b523b66..feebfd2e6 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -95,6 +95,7 @@ class InfectionMonkey: logger.info("Monkey is initializing...") self._singleton = SystemSingleton() self._opts = self._get_arguments(args) + # TODO: Change address to ip port to use list of servers self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.servers) self._control_client = ControlClient(self._opts.servers) # TODO Refactor the telemetry messengers to accept control client From 804bd4eadbfe63ca49d94cbab0eb4ca03fee2696 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 13:29:37 +0200 Subject: [PATCH 02/31] Agent: Modify find_server to accept list of servers --- monkey/infection_monkey/control.py | 93 +++++++++++++------ .../infection_monkey/test_control.py | 54 +++++++++++ 2 files changed, 117 insertions(+), 30 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 8d1e48a22..5ed3f24a5 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -1,20 +1,24 @@ import json import logging import platform +import socket from socket import gethostname -from typing import Mapping, Optional +from typing import Mapping, Optional, Sequence import requests from requests.exceptions import ConnectionError import infection_monkey.tunnel as tunnel from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT +from common.network.network_utils import address_to_ip_port from infection_monkey.config import GUID from infection_monkey.network.info import get_host_subnets, local_ips +from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE from infection_monkey.transport.http import HTTPConnectProxy from infection_monkey.transport.tcp import TcpProxy from infection_monkey.utils import agent_process from infection_monkey.utils.environment import is_windows_os +from infection_monkey.utils.threading import create_daemon_thread requests.packages.urllib3.disable_warnings() @@ -63,37 +67,66 @@ class ControlClient: timeout=MEDIUM_REQUEST_TIMEOUT, ) - def find_server(self, default_tunnel=None): - logger.debug(f"Trying to wake up with Monkey Island server: {self.server_address}") - if default_tunnel: - logger.debug("default_tunnel: %s" % (default_tunnel,)) + def find_server(self, servers: Sequence[str]): + logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") - try: - debug_message = "Trying to connect to server: %s" % self.server_address - if self.proxies: - debug_message += " through proxies: %s" % self.proxies - logger.debug(debug_message) - requests.get( # noqa: DUO123 - f"https://{self.server_address}/api?action=is-up", - verify=False, - proxies=self.proxies, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) - return True - except ConnectionError as exc: - logger.warning("Error connecting to control server %s: %s", self.server_address, exc) + while servers: + server = servers[0] - if self.proxies: - return False - else: - logger.info("Starting tunnel lookup...") - proxy_find = tunnel.find_tunnel(default=default_tunnel) - if proxy_find: - self.set_proxies(proxy_find) - return self.find_server() - else: - logger.info("No tunnel found") - return False + try: + debug_message = f"Trying to connect to server: {server}" + logger.debug(debug_message) + requests.get( # noqa: DUO123 + f"https://{server}/api?action=is-up", + verify=False, + timeout=MEDIUM_REQUEST_TIMEOUT, + ) + + # We remove the server that has been succesfull + servers.remove(server) + + # TODO: Check how we are going to set the server address that the ControlCLient + # is going to use + # self.server_address = server + + # If we have any other server we send them RELAY_CONTROL_MESSAGE + if servers: + for ss in servers: + t = create_daemon_thread( + target=ControlClient._send_relay_control_message, + name="SendControlRelayMessageThread", + args=(ss,), + ) + t.start() + t.join() + + return True + except ConnectionError as err: + logger.error(f"Unable to connect to server/relay {server}: {err}") + servers.remove(server) + except TimeoutError as err: + logger.error(f"Timed out while connecting to server/relay {server}: {err}") + servers.remove(server) + except Exception as err: + logger.error( + f"Exception encountered when trying to connect to server/relay {server}: {err}" + ) + servers.remove(server) + + return False + + @staticmethod + def _send_relay_control_message(server: str): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: + d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT) + + try: + address, port = address_to_ip_port(server) + d_socket.connect((address, int(port))) + d_socket.send(RELAY_CONTROL_MESSAGE) + logger.info(f"Control message was sent to the server/relay {server}") + except OSError as err: + logger.error(f"Error connecting to socket {server}: {err}") def set_proxies(self, proxy_find): """ diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index b90087ebf..6cd3df844 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -1,7 +1,31 @@ +from unittest.mock import MagicMock + import pytest +import requests from monkey.infection_monkey.control import ControlClient +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" + + +class MockConnectionError: + def __init__(self, *args, **kwargs): + raise requests.exceptions.ConnectionError + + +class RequestsGetArgument: + def __init__(self, *args, **kwargs): + if SERVER_1 in args[0]: + MockConnectionError() + + +@pytest.fixture +def servers(): + return [SERVER_1, SERVER_2, SERVER_3, SERVER_4] + @pytest.mark.parametrize( "is_windows_os,expected_proxy_string", @@ -14,3 +38,33 @@ def test_control_set_proxies(monkeypatch, is_windows_os, expected_proxy_string): control_client.set_proxies(("8.8.8.8", "45455")) assert control_client.proxies["https"] == expected_proxy_string + + +def test_control_find_server_any_exception(monkeypatch, servers): + monkeypatch.setattr("infection_monkey.control.requests.get", MockConnectionError) + + cc = ControlClient(servers) + + return_value = cc.find_server(servers) + + assert return_value is False + assert servers == [] + + +def test_control_find_server_socket(monkeypatch, servers): + mock_connect = MagicMock() + mock_send = MagicMock() + monkeypatch.setattr("infection_monkey.control.requests.get", RequestsGetArgument) + monkeypatch.setattr("infection_monkey.control.socket.socket.connect", mock_connect) + monkeypatch.setattr("infection_monkey.control.socket.socket.send", mock_send) + + cc = ControlClient(servers) + + return_value = cc.find_server(servers) + + assert len(servers) == 2 + assert return_value is True + assert mock_connect.call_count == 2 + assert mock_send.call_count == 2 + # TODO: be sure that connect is called with SERVER_3 and SERVER_4 + # assert mock_connect.call_args == SERVER_3 From a39917d9aa86807bb86ebcf99480bca13e09c814 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 17:57:47 +0530 Subject: [PATCH 03/31] UT: Fix test_control_find_server_socket to check call parameters --- .../unit_tests/infection_monkey/test_control.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 6cd3df844..2298589d6 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -1,8 +1,11 @@ +from unittest import mock from unittest.mock import MagicMock import pytest import requests +from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE +from monkey.common.network.network_utils import address_to_ip_port from monkey.infection_monkey.control import ControlClient SERVER_1 = "1.1.1.1:12312" @@ -65,6 +68,12 @@ def test_control_find_server_socket(monkeypatch, servers): assert len(servers) == 2 assert return_value is True assert mock_connect.call_count == 2 + + server_3_ip, server_3_port = address_to_ip_port(SERVER_3) + server_4_ip, server_4_port = address_to_ip_port(SERVER_4) + mock_connect.assert_has_calls( + [mock.call((server_3_ip, int(server_3_port))), mock.call((server_4_ip, int(server_4_port)))] + ) + assert mock_send.call_count == 2 - # TODO: be sure that connect is called with SERVER_3 and SERVER_4 - # assert mock_connect.call_args == SERVER_3 + mock_send.assert_called_with(RELAY_CONTROL_MESSAGE) From f847757a9a7e4a445364888abb23530c75b43405 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 18:00:08 +0530 Subject: [PATCH 04/31] UT: Rename test_control_find_server_any_exception -> test_control_find_server__no_available_relays --- monkey/tests/unit_tests/infection_monkey/test_control.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 2298589d6..181e38d9f 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -43,14 +43,12 @@ def test_control_set_proxies(monkeypatch, is_windows_os, expected_proxy_string): assert control_client.proxies["https"] == expected_proxy_string -def test_control_find_server_any_exception(monkeypatch, servers): +def test_control_find_server__no_available_relays(monkeypatch, servers): monkeypatch.setattr("infection_monkey.control.requests.get", MockConnectionError) cc = ControlClient(servers) - return_value = cc.find_server(servers) - - assert return_value is False + assert cc.find_server(servers) is False assert servers == [] From 47f838cf9f830e81c3cdcccaf14eb5f51311ff5e Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 18:01:56 +0530 Subject: [PATCH 05/31] UT: Rename test_control_find_server_socket -> test_control_find_server__control_message_sent_to_necessary_relays --- monkey/tests/unit_tests/infection_monkey/test_control.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 181e38d9f..7400e98ac 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -52,7 +52,7 @@ def test_control_find_server__no_available_relays(monkeypatch, servers): assert servers == [] -def test_control_find_server_socket(monkeypatch, servers): +def test_control_find_server__control_message_sent_to_necessary_relays(monkeypatch, servers): mock_connect = MagicMock() mock_send = MagicMock() monkeypatch.setattr("infection_monkey.control.requests.get", RequestsGetArgument) From 0239a1be6bcb9aad1356c93955d3a204621d08aa Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 18:07:22 +0530 Subject: [PATCH 06/31] UT: Simplify test logic in test_control_find_server__control_message_sent_to_necessary_relays --- monkey/tests/unit_tests/infection_monkey/test_control.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 7400e98ac..c5b8c0fb9 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -61,14 +61,12 @@ def test_control_find_server__control_message_sent_to_necessary_relays(monkeypat cc = ControlClient(servers) - return_value = cc.find_server(servers) - + assert cc.find_server(servers) is True assert len(servers) == 2 - assert return_value is True - assert mock_connect.call_count == 2 server_3_ip, server_3_port = address_to_ip_port(SERVER_3) server_4_ip, server_4_port = address_to_ip_port(SERVER_4) + assert mock_connect.call_count == 2 mock_connect.assert_has_calls( [mock.call((server_3_ip, int(server_3_port))), mock.call((server_4_ip, int(server_4_port)))] ) From cd91b3e42abbfdfe6f49d79f9c254cd8988ea5e4 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 18:12:24 +0530 Subject: [PATCH 07/31] UT: Rename RequestsGetArgument -> MockRequestsGetResponsePerServerArgument --- monkey/tests/unit_tests/infection_monkey/test_control.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index c5b8c0fb9..7131dc5ab 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -19,7 +19,7 @@ class MockConnectionError: raise requests.exceptions.ConnectionError -class RequestsGetArgument: +class MockRequestsGetResponsePerServerArgument: def __init__(self, *args, **kwargs): if SERVER_1 in args[0]: MockConnectionError() @@ -55,7 +55,9 @@ def test_control_find_server__no_available_relays(monkeypatch, servers): def test_control_find_server__control_message_sent_to_necessary_relays(monkeypatch, servers): mock_connect = MagicMock() mock_send = MagicMock() - monkeypatch.setattr("infection_monkey.control.requests.get", RequestsGetArgument) + monkeypatch.setattr( + "infection_monkey.control.requests.get", MockRequestsGetResponsePerServerArgument + ) monkeypatch.setattr("infection_monkey.control.socket.socket.connect", mock_connect) monkeypatch.setattr("infection_monkey.control.socket.socket.send", mock_send) From 02a919123b182920b181f9b4611be2453696ef0e Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 14:25:22 +0200 Subject: [PATCH 08/31] Agent: Remove thread join in find_server --- monkey/infection_monkey/control.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 5ed3f24a5..606f858fa 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -98,7 +98,6 @@ class ControlClient: args=(ss,), ) t.start() - t.join() return True except ConnectionError as err: From 178b296f755dc5b52627b83a9b29d50f2acb1e08 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 14:49:32 +0200 Subject: [PATCH 09/31] Agent: Use iterator in ControlClient.find_server --- monkey/infection_monkey/control.py | 32 ++++++++++-------------------- 1 file changed, 11 insertions(+), 21 deletions(-) diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 606f858fa..2d758d5cf 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -70,8 +70,9 @@ class ControlClient: def find_server(self, servers: Sequence[str]): logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") - while servers: - server = servers[0] + server_iterator = (s for s in servers) + + for server in server_iterator: try: debug_message = f"Trying to connect to server: {server}" @@ -82,37 +83,26 @@ class ControlClient: timeout=MEDIUM_REQUEST_TIMEOUT, ) - # We remove the server that has been succesfull - servers.remove(server) - + break # TODO: Check how we are going to set the server address that the ControlCLient # is going to use # self.server_address = server - - # If we have any other server we send them RELAY_CONTROL_MESSAGE - if servers: - for ss in servers: - t = create_daemon_thread( - target=ControlClient._send_relay_control_message, - name="SendControlRelayMessageThread", - args=(ss,), - ) - t.start() - - return True except ConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") - servers.remove(server) except TimeoutError as err: logger.error(f"Timed out while connecting to server/relay {server}: {err}") - servers.remove(server) except Exception as err: logger.error( f"Exception encountered when trying to connect to server/relay {server}: {err}" ) - servers.remove(server) - return False + for server in server_iterator: + t = create_daemon_thread( + target=ControlClient._send_relay_control_message, + name="SendControlRelayMessageThread", + args=(server,), + ) + t.start() @staticmethod def _send_relay_control_message(server: str): From 789d6b8441cb537a47ba12b20e1d958596cbd1fa Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 15:06:21 +0200 Subject: [PATCH 10/31] Agent: Move ControlClient.find_server in network/relay/utls.py --- monkey/infection_monkey/control.py | 57 +----------------- monkey/infection_monkey/monkey.py | 3 +- .../infection_monkey/network/relay/utils.py | 58 +++++++++++++++++++ 3 files changed, 61 insertions(+), 57 deletions(-) create mode 100644 monkey/infection_monkey/network/relay/utils.py diff --git a/monkey/infection_monkey/control.py b/monkey/infection_monkey/control.py index 2d758d5cf..3aedeb272 100644 --- a/monkey/infection_monkey/control.py +++ b/monkey/infection_monkey/control.py @@ -1,24 +1,19 @@ import json import logging import platform -import socket from socket import gethostname -from typing import Mapping, Optional, Sequence +from typing import Mapping, Optional import requests -from requests.exceptions import ConnectionError import infection_monkey.tunnel as tunnel from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT -from common.network.network_utils import address_to_ip_port from infection_monkey.config import GUID from infection_monkey.network.info import get_host_subnets, local_ips -from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE from infection_monkey.transport.http import HTTPConnectProxy from infection_monkey.transport.tcp import TcpProxy from infection_monkey.utils import agent_process from infection_monkey.utils.environment import is_windows_os -from infection_monkey.utils.threading import create_daemon_thread requests.packages.urllib3.disable_warnings() @@ -67,56 +62,6 @@ class ControlClient: timeout=MEDIUM_REQUEST_TIMEOUT, ) - def find_server(self, servers: Sequence[str]): - logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") - - server_iterator = (s for s in servers) - - for server in server_iterator: - - try: - debug_message = f"Trying to connect to server: {server}" - logger.debug(debug_message) - requests.get( # noqa: DUO123 - f"https://{server}/api?action=is-up", - verify=False, - timeout=MEDIUM_REQUEST_TIMEOUT, - ) - - break - # TODO: Check how we are going to set the server address that the ControlCLient - # is going to use - # self.server_address = server - except ConnectionError as err: - logger.error(f"Unable to connect to server/relay {server}: {err}") - except TimeoutError as err: - logger.error(f"Timed out while connecting to server/relay {server}: {err}") - except Exception as err: - logger.error( - f"Exception encountered when trying to connect to server/relay {server}: {err}" - ) - - for server in server_iterator: - t = create_daemon_thread( - target=ControlClient._send_relay_control_message, - name="SendControlRelayMessageThread", - args=(server,), - ) - t.start() - - @staticmethod - def _send_relay_control_message(server: str): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: - d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT) - - try: - address, port = address_to_ip_port(server) - d_socket.connect((address, int(port))) - d_socket.send(RELAY_CONTROL_MESSAGE) - logger.info(f"Control message was sent to the server/relay {server}") - except OSError as err: - logger.error(f"Error connecting to socket {server}: {err}") - def set_proxies(self, proxy_find): """ Note: The proxy schema changes between different versions of requests and urllib3, diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index feebfd2e6..1af0266bb 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -42,6 +42,7 @@ from infection_monkey.master.control_channel import ControlChannel from infection_monkey.model import VictimHostFactory from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_network_interfaces +from infection_monkey.network.relay.utils import find_server from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter @@ -162,7 +163,7 @@ class InfectionMonkey: self._control_client.wakeup(parent=self._opts.parent) def _current_server_is_set(self) -> bool: - if self._control_client.find_server(default_tunnel=self._opts.servers): + if find_server(servers=self._opts.servers): return True return False diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py new file mode 100644 index 000000000..56462dc33 --- /dev/null +++ b/monkey/infection_monkey/network/relay/utils.py @@ -0,0 +1,58 @@ +import logging +import socket +from typing import Sequence + +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 +from infection_monkey.utils.threading import create_daemon_thread + +logger = logging.getLogger(__name__) + + +def find_server(self, servers: Sequence[str]): + logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") + + server_iterator = (s for s in servers) + for server in server_iterator: + try: + debug_message = f"Trying to connect to server: {server}" + logger.debug(debug_message) + requests.get( # noqa: DUO123 + f"https://{server}/api?action=is-up", + verify=False, + timeout=MEDIUM_REQUEST_TIMEOUT, + ) + + break + except requests.exceptions.ConnectionError as err: + logger.error(f"Unable to connect to server/relay {server}: {err}") + except TimeoutError as err: + logger.error(f"Timed out while connecting to server/relay {server}: {err}") + except Exception as err: + logger.error( + f"Exception encountered when trying to connect to server/relay {server}: {err}" + ) + + for server in server_iterator: + t = create_daemon_thread( + target=_send_relay_control_message, + name="SendControlRelayMessageThread", + args=(server,), + ) + t.start() + + +def _send_relay_control_message(server: str): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: + d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT) + + try: + address, port = address_to_ip_port(server) + d_socket.connect((address, int(port))) + d_socket.send(RELAY_CONTROL_MESSAGE) + logger.info(f"Control message was sent to the server/relay {server}") + except OSError as err: + logger.error(f"Error connecting to socket {server}: {err}") From 18659b654fdea601574573a68fb85829cab42dcd Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 18:57:21 +0530 Subject: [PATCH 11/31] Agent: Return server to connect to or None from `find_server` --- monkey/infection_monkey/network/relay/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 56462dc33..66f0b1d1e 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,6 +1,6 @@ import logging import socket -from typing import Sequence +from typing import Optional, Sequence import requests @@ -12,7 +12,9 @@ from infection_monkey.utils.threading import create_daemon_thread logger = logging.getLogger(__name__) -def find_server(self, servers: Sequence[str]): +def find_server(self, servers: Sequence[str]) -> Optional[str]: + server_found = None + logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") server_iterator = (s for s in servers) @@ -26,6 +28,8 @@ def find_server(self, servers: Sequence[str]): timeout=MEDIUM_REQUEST_TIMEOUT, ) + server_found = server + break except requests.exceptions.ConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") @@ -44,6 +48,8 @@ def find_server(self, servers: Sequence[str]): ) t.start() + return server_found + def _send_relay_control_message(server: str): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: From 94dc8cf37785b3b9d1a8ca50d160df98d25a1d82 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 19:00:51 +0530 Subject: [PATCH 12/31] Agent: Use `find_servers` to pass valid server to `ControlClient` --- monkey/infection_monkey/monkey.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 1af0266bb..2a70b763c 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -96,9 +96,11 @@ class InfectionMonkey: logger.info("Monkey is initializing...") self._singleton = SystemSingleton() self._opts = self._get_arguments(args) - # TODO: Change address to ip port to use list of servers - self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self._opts.servers) - self._control_client = ControlClient(self._opts.servers) + + server = find_server(self._opts.servers) + self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) + self._control_client = ControlClient(server_address=server) + # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object ControlClient.control_client_object = self._control_client From c6c6cf1e7926fb36737ca1750679e62f13f38397 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Wed, 7 Sep 2022 19:02:14 +0530 Subject: [PATCH 13/31] Agent: Add TODO about variable naming in `InfectionMonkey` --- monkey/infection_monkey/monkey.py | 1 + 1 file changed, 1 insertion(+) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 2a70b763c..596cd0d7f 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -97,6 +97,7 @@ class InfectionMonkey: self._singleton = SystemSingleton() self._opts = self._get_arguments(args) + # TODO: Revisit variable names server = find_server(self._opts.servers) self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) self._control_client = ControlClient(server_address=server) From bb2b4aaf6c2b58330e1ef44254e4e99e742ecc6d Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 16:45:47 +0200 Subject: [PATCH 14/31] Agent: Separate responsibilites in network.relay.utils.find_server --- .../infection_monkey/network/relay/utils.py | 19 +++++----- .../network/relay/test_utils.py | 35 +++++++++++++++++++ 2 files changed, 45 insertions(+), 9 deletions(-) create mode 100644 monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 66f0b1d1e..8d3a76277 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -1,6 +1,6 @@ import logging import socket -from typing import Optional, Sequence +from typing import Iterable, Optional import requests @@ -12,13 +12,12 @@ from infection_monkey.utils.threading import create_daemon_thread logger = logging.getLogger(__name__) -def find_server(self, servers: Sequence[str]) -> Optional[str]: +def find_server(servers: Iterable[str]) -> Optional[str]: server_found = None logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") - server_iterator = (s for s in servers) - for server in server_iterator: + for server in servers: try: debug_message = f"Trying to connect to server: {server}" logger.debug(debug_message) @@ -40,18 +39,20 @@ def find_server(self, servers: Sequence[str]) -> Optional[str]: f"Exception encountered when trying to connect to server/relay {server}: {err}" ) - for server in server_iterator: + return server_found + + +def send_relay_control_message(servers: Iterable[str]): + for server in servers: t = create_daemon_thread( - target=_send_relay_control_message, + target=_open_socket_to_server, name="SendControlRelayMessageThread", args=(server,), ) t.start() - return server_found - -def _send_relay_control_message(server: str): +def _open_socket_to_server(server: str): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT) 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 new file mode 100644 index 000000000..15c4f46da --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_utils.py @@ -0,0 +1,35 @@ +import pytest +import requests + +from infection_monkey.network.relay.utils import find_server + +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" + + +class MockConnectionError: + def __init__(self, *args, **kwargs): + raise requests.exceptions.ConnectionError + + +class MockRequestsGetResponsePerServerArgument: + def __init__(self, *args, **kwargs): + if SERVER_1 in args[0]: + MockConnectionError() + + +@pytest.fixture +def servers(): + return [SERVER_1, SERVER_2, SERVER_3, SERVER_4] + + +@pytest.mark.parametrize( + "mock_requests_get, expected", + [(MockConnectionError, None), (MockRequestsGetResponsePerServerArgument, SERVER_2)], +) +def test_find_server__no_available_relays(monkeypatch, servers, mock_requests_get, expected): + monkeypatch.setattr("infection_monkey.control.requests.get", mock_requests_get) + + assert find_server(servers) is expected From 65226d5a9c4768547cd41ef64a9ca40a3076808f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 16:46:44 +0200 Subject: [PATCH 15/31] UT: Move test_find_server under network.relay.utils --- .../infection_monkey/test_control.py | 61 ------------------- 1 file changed, 61 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/test_control.py b/monkey/tests/unit_tests/infection_monkey/test_control.py index 7131dc5ab..b90087ebf 100644 --- a/monkey/tests/unit_tests/infection_monkey/test_control.py +++ b/monkey/tests/unit_tests/infection_monkey/test_control.py @@ -1,34 +1,7 @@ -from unittest import mock -from unittest.mock import MagicMock - import pytest -import requests -from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE -from monkey.common.network.network_utils import address_to_ip_port from monkey.infection_monkey.control import ControlClient -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" - - -class MockConnectionError: - def __init__(self, *args, **kwargs): - raise requests.exceptions.ConnectionError - - -class MockRequestsGetResponsePerServerArgument: - def __init__(self, *args, **kwargs): - if SERVER_1 in args[0]: - MockConnectionError() - - -@pytest.fixture -def servers(): - return [SERVER_1, SERVER_2, SERVER_3, SERVER_4] - @pytest.mark.parametrize( "is_windows_os,expected_proxy_string", @@ -41,37 +14,3 @@ def test_control_set_proxies(monkeypatch, is_windows_os, expected_proxy_string): control_client.set_proxies(("8.8.8.8", "45455")) assert control_client.proxies["https"] == expected_proxy_string - - -def test_control_find_server__no_available_relays(monkeypatch, servers): - monkeypatch.setattr("infection_monkey.control.requests.get", MockConnectionError) - - cc = ControlClient(servers) - - assert cc.find_server(servers) is False - assert servers == [] - - -def test_control_find_server__control_message_sent_to_necessary_relays(monkeypatch, servers): - mock_connect = MagicMock() - mock_send = MagicMock() - monkeypatch.setattr( - "infection_monkey.control.requests.get", MockRequestsGetResponsePerServerArgument - ) - monkeypatch.setattr("infection_monkey.control.socket.socket.connect", mock_connect) - monkeypatch.setattr("infection_monkey.control.socket.socket.send", mock_send) - - cc = ControlClient(servers) - - assert cc.find_server(servers) is True - assert len(servers) == 2 - - server_3_ip, server_3_port = address_to_ip_port(SERVER_3) - server_4_ip, server_4_port = address_to_ip_port(SERVER_4) - assert mock_connect.call_count == 2 - mock_connect.assert_has_calls( - [mock.call((server_3_ip, int(server_3_port))), mock.call((server_4_ip, int(server_4_port)))] - ) - - assert mock_send.call_count == 2 - mock_send.assert_called_with(RELAY_CONTROL_MESSAGE) From e539495545ff554c7787ec027c8a91b97f7b0f94 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 16:47:47 +0200 Subject: [PATCH 16/31] Agent: Find server and send control relay message to all other servers --- monkey/infection_monkey/monkey.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 596cd0d7f..539d5c61d 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -42,7 +42,7 @@ from infection_monkey.master.control_channel import ControlChannel from infection_monkey.model import VictimHostFactory from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_network_interfaces -from infection_monkey.network.relay.utils import find_server +from infection_monkey.network.relay.utils import find_server, send_relay_control_message from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter @@ -98,9 +98,9 @@ class InfectionMonkey: self._opts = self._get_arguments(args) # TODO: Revisit variable names - server = find_server(self._opts.servers) - self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) - self._control_client = ControlClient(server_address=server) + self._get_server() + self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self.server) + self._control_client = ControlClient(server_address=self.server) # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object @@ -122,6 +122,11 @@ class InfectionMonkey: return opts + def _get_server(self): + servers_iterator = (s for s in self._opts.servers) + self.server = find_server(servers_iterator) + send_relay_control_message(servers_iterator) + @staticmethod def _log_arguments(args): arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()]) @@ -159,14 +164,13 @@ class InfectionMonkey: logger.debug(f"Default server set to: {self._control_client.server_address}") else: raise Exception( - f"Failed to connect to the island via " - f"any known server address: {self._opts.servers}" + f"Failed to connect to the island via " f"any known server address: {self.server}" ) self._control_client.wakeup(parent=self._opts.parent) def _current_server_is_set(self) -> bool: - if find_server(servers=self._opts.servers): + if self.server: return True return False From 60f9aa6a4e479937dbb8fe1e8e99f733dddd8d1f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 19:15:25 +0200 Subject: [PATCH 17/31] Agent: Rework relay.utils.find_server a bit --- monkey/infection_monkey/network/relay/utils.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 8d3a76277..459900e0d 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -13,23 +13,19 @@ logger = logging.getLogger(__name__) def find_server(servers: Iterable[str]) -> Optional[str]: - server_found = None - logger.debug(f"Trying to wake up with servers: {', '.join(servers)}") for server in servers: + logger.debug(f"Trying to connect to server: {server}") + try: - debug_message = f"Trying to connect to server: {server}" - logger.debug(debug_message) requests.get( # noqa: DUO123 f"https://{server}/api?action=is-up", verify=False, timeout=MEDIUM_REQUEST_TIMEOUT, ) - server_found = server - - break + return server except requests.exceptions.ConnectionError as err: logger.error(f"Unable to connect to server/relay {server}: {err}") except TimeoutError as err: @@ -39,7 +35,7 @@ def find_server(servers: Iterable[str]) -> Optional[str]: f"Exception encountered when trying to connect to server/relay {server}: {err}" ) - return server_found + return None def send_relay_control_message(servers: Iterable[str]): From 20172230f180314ce07e12dc06a0e8b09f81975f Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 19:39:14 +0200 Subject: [PATCH 18/31] Agent: Rework send_control_relay_message a bit --- monkey/infection_monkey/network/relay/utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 459900e0d..c6f479f8a 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -52,10 +52,13 @@ def _open_socket_to_server(server: str): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT) + ip, port = address_to_ip_port(server) + logger.info(f"Control message was sent to the server/relay {server}") + try: - address, port = address_to_ip_port(server) - d_socket.connect((address, int(port))) + d_socket.connect((ip, int(port))) d_socket.send(RELAY_CONTROL_MESSAGE) - logger.info(f"Control message was sent to the server/relay {server}") except OSError as err: logger.error(f"Error connecting to socket {server}: {err}") + except TimeoutError as err: + logger.error(f"Timed out while connecting to socket {server}: {err}") From b89ba06fd16cfa4506974aa90a687d80c7137727 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 19:47:14 +0200 Subject: [PATCH 19/31] UT: Rename test_find_server__no_available_relays to test_find_server --- .../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 15c4f46da..bf4d70056 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 @@ -29,7 +29,7 @@ def servers(): "mock_requests_get, expected", [(MockConnectionError, None), (MockRequestsGetResponsePerServerArgument, SERVER_2)], ) -def test_find_server__no_available_relays(monkeypatch, servers, mock_requests_get, expected): +def test_find_server(monkeypatch, servers, mock_requests_get, expected): monkeypatch.setattr("infection_monkey.control.requests.get", mock_requests_get) assert find_server(servers) is expected From 6ee15e22b83c8fa7c69a68c4ad18d601ac069029 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Wed, 7 Sep 2022 19:53:54 +0200 Subject: [PATCH 20/31] Agent: Rework call of find_server in monkey.py --- monkey/infection_monkey/monkey.py | 36 ++++++++++++------------------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 539d5c61d..a7dfe4aa8 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -98,9 +98,9 @@ class InfectionMonkey: self._opts = self._get_arguments(args) # TODO: Revisit variable names - self._get_server() - self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(self.server) - self._control_client = ControlClient(server_address=self.server) + server = self._get_server() + self._cmd_island_ip, self._cmd_island_port = address_to_ip_port(server) + self._control_client = ControlClient(server_address=server) # TODO Refactor the telemetry messengers to accept control client # and remove control_client_object @@ -124,9 +124,18 @@ class InfectionMonkey: def _get_server(self): servers_iterator = (s for s in self._opts.servers) - self.server = find_server(servers_iterator) + server = find_server(servers_iterator) + if server: + logger.debug(f"Default server set to: {server}") + else: + # TODO: Exit here + raise Exception( + f"Failed to connect to the island via any known servers: {self._opts.servers}" + ) send_relay_control_message(servers_iterator) + return server + @staticmethod def _log_arguments(args): arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()]) @@ -140,7 +149,7 @@ class InfectionMonkey: logger.info("Agent is starting...") logger.info(f"Agent GUID: {GUID}") - self._connect_to_island() + self._control_client.wakeup(parent=self._opts.parent) # TODO: Reevaluate who is responsible to send this information if is_windows_os(): @@ -158,23 +167,6 @@ class InfectionMonkey: self._setup() self._master.start() - def _connect_to_island(self): - # Sets island's IP and port for monkey to communicate to - if self._current_server_is_set(): - logger.debug(f"Default server set to: {self._control_client.server_address}") - else: - raise Exception( - f"Failed to connect to the island via " f"any known server address: {self.server}" - ) - - self._control_client.wakeup(parent=self._opts.parent) - - def _current_server_is_set(self) -> bool: - if self.server: - return True - - return False - def _setup(self): logger.debug("Starting the setup phase.") From fb1554840ac5f87e8e21cc26040e63b3118250bf Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:22:13 +0530 Subject: [PATCH 21/31] Agent: Fix log message in InfectionMonkey._get_server() --- 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 a7dfe4aa8..efa3046b0 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -126,7 +126,7 @@ class InfectionMonkey: servers_iterator = (s for s in self._opts.servers) server = find_server(servers_iterator) if server: - logger.debug(f"Default server set to: {server}") + logger.info(f"Successfully connected to the island via {server}") else: # TODO: Exit here raise Exception( From f436bf7b8c757d9e561310e1ca885b0ec921ce45 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:24:11 +0530 Subject: [PATCH 22/31] Agent: Remove irrelevant comment from InfectionMonkey._get_server() --- monkey/infection_monkey/monkey.py | 1 - 1 file changed, 1 deletion(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index efa3046b0..4a8685aae 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -128,7 +128,6 @@ class InfectionMonkey: if server: logger.info(f"Successfully connected to the island via {server}") else: - # TODO: Exit here raise Exception( f"Failed to connect to the island via any known servers: {self._opts.servers}" ) From 78d32053a219551006743737e27b71dfbcd8f3ad Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:30:56 +0530 Subject: [PATCH 23/31] Agent: Rename send_relay_control_message -> send_remove_from_waitlist_control_message_to_relays --- monkey/infection_monkey/monkey.py | 7 +++++-- monkey/infection_monkey/network/relay/utils.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 4a8685aae..6d8e6bac6 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -42,7 +42,10 @@ from infection_monkey.master.control_channel import ControlChannel from infection_monkey.model import VictimHostFactory from infection_monkey.network.firewall import app as firewall from infection_monkey.network.info import get_network_interfaces -from infection_monkey.network.relay.utils import find_server, send_relay_control_message +from infection_monkey.network.relay.utils import ( + find_server, + send_remove_from_waitlist_control_message_to_relays, +) from infection_monkey.network_scanning.elasticsearch_fingerprinter import ElasticSearchFingerprinter from infection_monkey.network_scanning.http_fingerprinter import HTTPFingerprinter from infection_monkey.network_scanning.mssql_fingerprinter import MSSQLFingerprinter @@ -131,7 +134,7 @@ class InfectionMonkey: raise Exception( f"Failed to connect to the island via any known servers: {self._opts.servers}" ) - send_relay_control_message(servers_iterator) + send_remove_from_waitlist_control_message_to_relays(servers_iterator) return server diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index c6f479f8a..91bbf75fc 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -38,7 +38,7 @@ def find_server(servers: Iterable[str]) -> Optional[str]: return None -def send_relay_control_message(servers: Iterable[str]): +def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): for server in servers: t = create_daemon_thread( target=_open_socket_to_server, From 6bfe6bc79d3d480bab39af6f54aeda2442e6ace5 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:31:56 +0530 Subject: [PATCH 24/31] Common: Rename SendControlRelayMessageThread -> SendRemoveFromWaitlistControlMessageToRelaysThread --- 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 91bbf75fc..956945068 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -42,7 +42,7 @@ def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): for server in servers: t = create_daemon_thread( target=_open_socket_to_server, - name="SendControlRelayMessageThread", + name="SendRemoveFromWaitlistControlMessageToRelaysThread", args=(server,), ) t.start() From 7661027c6c50b42e2f2b340f30b4a31bc55791b4 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:37:25 +0530 Subject: [PATCH 25/31] Agent: Don't catch TimeoutError in _open_socket_to_server() since OSError is already being caught --- monkey/infection_monkey/network/relay/utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 956945068..5923002b6 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -60,5 +60,3 @@ def _open_socket_to_server(server: str): d_socket.send(RELAY_CONTROL_MESSAGE) except OSError as err: logger.error(f"Error connecting to socket {server}: {err}") - except TimeoutError as err: - logger.error(f"Timed out while connecting to socket {server}: {err}") From cb4af415c1773a8522ffa2c4429d88a1988692e6 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:39:45 +0530 Subject: [PATCH 26/31] Agent: Rename _open_socket_to_server -> _send_remove_from_waitlist_control_message_to_relay --- 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 5923002b6..333eafd6d 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -41,14 +41,14 @@ def find_server(servers: Iterable[str]) -> Optional[str]: def send_remove_from_waitlist_control_message_to_relays(servers: Iterable[str]): for server in servers: t = create_daemon_thread( - target=_open_socket_to_server, + target=_send_remove_from_waitlist_control_message_to_relay, name="SendRemoveFromWaitlistControlMessageToRelaysThread", args=(server,), ) t.start() -def _open_socket_to_server(server: str): +def _send_remove_from_waitlist_control_message_to_relay(server: str): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as d_socket: d_socket.settimeout(MEDIUM_REQUEST_TIMEOUT) From aa1c31efb2560a99c9d8c883112b93843774b58a Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:42:11 +0530 Subject: [PATCH 27/31] Agent: Rename RELAY_CONTROL_MESSAGE -> RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST --- monkey/infection_monkey/network/relay/__init__.py | 5 ++++- .../network/relay/relay_connection_handler.py | 4 ++-- monkey/infection_monkey/network/relay/utils.py | 4 ++-- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/monkey/infection_monkey/network/relay/__init__.py b/monkey/infection_monkey/network/relay/__init__.py index ad5d325c5..61728b643 100644 --- a/monkey/infection_monkey/network/relay/__init__.py +++ b/monkey/infection_monkey/network/relay/__init__.py @@ -1,4 +1,7 @@ -from .relay_connection_handler import RelayConnectionHandler, RELAY_CONTROL_MESSAGE +from .relay_connection_handler import ( + RelayConnectionHandler, + RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST, +) from .relay_user_handler import RelayUser, RelayUserHandler from .sockets_pipe import SocketsPipe from .tcp_connection_handler import TCPConnectionHandler diff --git a/monkey/infection_monkey/network/relay/relay_connection_handler.py b/monkey/infection_monkey/network/relay/relay_connection_handler.py index 3d7755fb9..4b4475e52 100644 --- a/monkey/infection_monkey/network/relay/relay_connection_handler.py +++ b/monkey/infection_monkey/network/relay/relay_connection_handler.py @@ -4,7 +4,7 @@ from ipaddress import IPv4Address from .relay_user_handler import RelayUserHandler from .tcp_pipe_spawner import TCPPipeSpawner -RELAY_CONTROL_MESSAGE = b"infection-monkey-relay-control-message: -" +RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST = b"infection-monkey-relay-control-message: -" class RelayConnectionHandler: @@ -25,7 +25,7 @@ class RelayConnectionHandler: control_message = sock.recv(socket.MSG_PEEK) - if control_message.startswith(RELAY_CONTROL_MESSAGE): + if control_message.startswith(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST): self._relay_user_handler.disconnect_user(addr) else: self._relay_user_handler.add_relay_user(addr) diff --git a/monkey/infection_monkey/network/relay/utils.py b/monkey/infection_monkey/network/relay/utils.py index 333eafd6d..096500c5b 100644 --- a/monkey/infection_monkey/network/relay/utils.py +++ b/monkey/infection_monkey/network/relay/utils.py @@ -6,7 +6,7 @@ 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 +from infection_monkey.network.relay import RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST from infection_monkey.utils.threading import create_daemon_thread logger = logging.getLogger(__name__) @@ -57,6 +57,6 @@ def _send_remove_from_waitlist_control_message_to_relay(server: str): try: d_socket.connect((ip, int(port))) - d_socket.send(RELAY_CONTROL_MESSAGE) + d_socket.send(RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST) except OSError as err: logger.error(f"Error connecting to socket {server}: {err}") From 451d2d069473736c0cb2d5068b9e1dcf4052d1d7 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 13:43:38 +0530 Subject: [PATCH 28/31] UT: Use RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST in place of RELAY_CONTROL_MESSAGE --- .../network/relay/test_relay_connection_handler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py index 92aff10fc..48f8b6bf5 100644 --- a/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py +++ b/monkey/tests/unit_tests/infection_monkey/network/relay/test_relay_connection_handler.py @@ -5,7 +5,7 @@ from unittest.mock import MagicMock import pytest from monkey.infection_monkey.network.relay import ( - RELAY_CONTROL_MESSAGE, + RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST, RelayConnectionHandler, RelayUserHandler, TCPPipeSpawner, @@ -27,7 +27,7 @@ def relay_user_handler(): @pytest.fixture def close_socket(): sock = MagicMock(spec=socket.socket) - sock.recv.return_value = RELAY_CONTROL_MESSAGE + sock.recv.return_value = RELAY_CONTROL_MESSAGE_REMOVE_FROM_WAITLIST sock.getpeername.return_value = (USER_ADDRESS, 12345) return sock From 45d1cc78c1814e95ed74b780f6ed2e665a6545a7 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 14:38:19 +0530 Subject: [PATCH 29/31] UT: Use requests_mock in test_find_server --- .../network/relay/test_utils.py | 30 +++++++------------ 1 file changed, 11 insertions(+), 19 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 bf4d70056..cae3cacbd 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,6 @@ import pytest import requests +import requests_mock from infection_monkey.network.relay.utils import find_server @@ -9,27 +10,18 @@ SERVER_3 = "3.3.3.3:3142" SERVER_4 = "4.4.4.4:5000" -class MockConnectionError: - def __init__(self, *args, **kwargs): - raise requests.exceptions.ConnectionError - - -class MockRequestsGetResponsePerServerArgument: - def __init__(self, *args, **kwargs): - if SERVER_1 in args[0]: - MockConnectionError() - - -@pytest.fixture -def servers(): - return [SERVER_1, SERVER_2, SERVER_3, SERVER_4] +servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( - "mock_requests_get, expected", - [(MockConnectionError, None), (MockRequestsGetResponsePerServerArgument, SERVER_2)], + "expected_server,connection_error_servers,do_nothing_servers", + [(None, servers, []), (SERVER_2, [SERVER_1], [SERVER_2, SERVER_3, SERVER_4])], ) -def test_find_server(monkeypatch, servers, mock_requests_get, expected): - monkeypatch.setattr("infection_monkey.control.requests.get", mock_requests_get) +def test_find_server(expected_server, connection_error_servers, do_nothing_servers): + with requests_mock.Mocker() as mock: + for server in connection_error_servers: + mock.get(f"https://{server}/api?action=is-up", exc=requests.exceptions.ConnectionError) + for server in do_nothing_servers: + mock.get(f"https://{server}/api?action=is-up", text="") - assert find_server(servers) is expected + assert find_server(servers) is expected_server From fac179bbda165bb1a62527be1aec111f0ec440d9 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 14:48:03 +0530 Subject: [PATCH 30/31] UT: Simplify test logic in test_find_server() --- .../network/relay/test_utils.py | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 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 cae3cacbd..fa63d34cb 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 @@ -14,14 +14,31 @@ servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( - "expected_server,connection_error_servers,do_nothing_servers", - [(None, servers, []), (SERVER_2, [SERVER_1], [SERVER_2, SERVER_3, SERVER_4])], + "expected_server,server_response_pairs", + [ + ( + None, + [ + (SERVER_1, {"exc": requests.exceptions.ConnectionError}), + (SERVER_2, {"exc": requests.exceptions.ConnectionError}), + (SERVER_3, {"exc": requests.exceptions.ConnectionError}), + (SERVER_4, {"exc": requests.exceptions.ConnectionError}), + ], + ), + ( + SERVER_2, + [ + (SERVER_1, {"exc": requests.exceptions.ConnectionError}), + (SERVER_2, {"text": ""}), + (SERVER_3, {"text": ""}), + (SERVER_4, {"text": ""}), + ], + ), + ], ) -def test_find_server(expected_server, connection_error_servers, do_nothing_servers): +def test_find_server(expected_server, server_response_pairs): with requests_mock.Mocker() as mock: - for server in connection_error_servers: - mock.get(f"https://{server}/api?action=is-up", exc=requests.exceptions.ConnectionError) - for server in do_nothing_servers: - mock.get(f"https://{server}/api?action=is-up", text="") + for server, response in server_response_pairs: + mock.get(f"https://{server}/api?action=is-up", **response) assert find_server(servers) is expected_server From e1759a7906d243e1a8a92bdb6474688d8d1693e0 Mon Sep 17 00:00:00 2001 From: Shreya Malviya Date: Thu, 8 Sep 2022 14:53:01 +0530 Subject: [PATCH 31/31] UT: Simplify parametrize logic in test_find_server() --- .../network/relay/test_utils.py | 18 +++--------------- 1 file changed, 3 insertions(+), 15 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 fa63d34cb..1acd08012 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 @@ -16,23 +16,11 @@ servers = [SERVER_1, SERVER_2, SERVER_3, SERVER_4] @pytest.mark.parametrize( "expected_server,server_response_pairs", [ - ( - None, - [ - (SERVER_1, {"exc": requests.exceptions.ConnectionError}), - (SERVER_2, {"exc": requests.exceptions.ConnectionError}), - (SERVER_3, {"exc": requests.exceptions.ConnectionError}), - (SERVER_4, {"exc": requests.exceptions.ConnectionError}), - ], - ), + (None, [(server, {"exc": requests.exceptions.ConnectionError}) for server in servers]), ( SERVER_2, - [ - (SERVER_1, {"exc": requests.exceptions.ConnectionError}), - (SERVER_2, {"text": ""}), - (SERVER_3, {"text": ""}), - (SERVER_4, {"text": ""}), - ], + [(SERVER_1, {"exc": requests.exceptions.ConnectionError})] + + [(server, {"text": ""}) for server in servers[1:]], ), ], )