From 4982999b999c19ecb1108c35ce7342fb14a1a088 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 23 Sep 2022 17:54:49 +0000 Subject: [PATCH 1/3] Common: Add function to parse SocketAddress --- monkey/common/types/__init__.py | 1 + monkey/common/{ => types}/types.py | 0 monkey/common/types/utils/__init__.py | 1 + monkey/common/types/utils/socket_address.py | 18 ++++++++++ .../common/types/utils/test_socket_address.py | 35 +++++++++++++++++++ 5 files changed, 55 insertions(+) create mode 100644 monkey/common/types/__init__.py rename monkey/common/{ => types}/types.py (100%) create mode 100644 monkey/common/types/utils/__init__.py create mode 100644 monkey/common/types/utils/socket_address.py create mode 100644 monkey/tests/unit_tests/common/types/utils/test_socket_address.py diff --git a/monkey/common/types/__init__.py b/monkey/common/types/__init__.py new file mode 100644 index 000000000..97bfe1c80 --- /dev/null +++ b/monkey/common/types/__init__.py @@ -0,0 +1 @@ +from .types import AgentID, HardwareID, MachineID, SocketAddress diff --git a/monkey/common/types.py b/monkey/common/types/types.py similarity index 100% rename from monkey/common/types.py rename to monkey/common/types/types.py diff --git a/monkey/common/types/utils/__init__.py b/monkey/common/types/utils/__init__.py new file mode 100644 index 000000000..7d7440ee3 --- /dev/null +++ b/monkey/common/types/utils/__init__.py @@ -0,0 +1 @@ +from .socket_address import socketaddress_from_string diff --git a/monkey/common/types/utils/socket_address.py b/monkey/common/types/utils/socket_address.py new file mode 100644 index 000000000..77418d4be --- /dev/null +++ b/monkey/common/types/utils/socket_address.py @@ -0,0 +1,18 @@ +from ipaddress import IPv4Address + +from common.network.network_utils import address_to_ip_port +from common.types import SocketAddress + + +def socketaddress_from_string(address_str: str) -> SocketAddress: + """ + Parse a SocketAddress object from a string + + :param address_str: A string of ip:port + :raises ValueError: If the string is not a valid ip:port + :return: SocketAddress with the IP and port + """ + ip, port = address_to_ip_port(address_str) + if port is None: + raise ValueError("SocketAddress requires a port") + return SocketAddress(ip=IPv4Address(ip), port=int(port)) diff --git a/monkey/tests/unit_tests/common/types/utils/test_socket_address.py b/monkey/tests/unit_tests/common/types/utils/test_socket_address.py new file mode 100644 index 000000000..a73fe8c6f --- /dev/null +++ b/monkey/tests/unit_tests/common/types/utils/test_socket_address.py @@ -0,0 +1,35 @@ +import pytest + +from common.types import SocketAddress +from common.types.utils import socketaddress_from_string + +GOOD_IP = "192.168.1.1" +BAD_IP = "192.168.1.999" +GOOD_PORT = 1234 +BAD_PORT = 99999 + + +def test_socketaddress_from_string(): + expected = SocketAddress(ip=GOOD_IP, port=GOOD_PORT) + + address = socketaddress_from_string(f"{GOOD_IP}:{GOOD_PORT}") + + assert address == expected + + +@pytest.mark.parametrize( + "bad_address", + [ + "not an address", + ":", + GOOD_IP, + str(GOOD_PORT), + f"{GOOD_IP}:", + f":{GOOD_PORT}", + f"{BAD_IP}:{GOOD_PORT}", + f"{GOOD_IP}:{BAD_PORT}", + ], +) +def test_socketaddress_from_string__raises(bad_address: str): + with pytest.raises(ValueError): + socketaddress_from_string(bad_address) From 53a9c622454f8fbce01c6dc68da3c30975afcce6 Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Fri, 23 Sep 2022 17:57:53 +0000 Subject: [PATCH 2/3] Agent: Parse --servers to SocketAddress --- monkey/infection_monkey/monkey.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index b5c20fd76..1895211ea 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -21,6 +21,7 @@ from common.network.network_utils import ( get_my_ip_addresses, get_network_interfaces, ) +from common.types.utils import socketaddress_from_string from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version @@ -113,6 +114,7 @@ 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() @@ -138,7 +140,11 @@ class InfectionMonkey: def _get_arguments(args): arg_parser = argparse.ArgumentParser() arg_parser.add_argument("-p", "--parent") - arg_parser.add_argument("-s", "--servers", type=lambda arg: arg.strip().split(",")) + arg_parser.add_argument( + "-s", + "--servers", + type=lambda arg: [socketaddress_from_string(s) for s in arg.strip().split(",")], + ) arg_parser.add_argument("-d", "--depth", type=positive_int, default=0) opts = arg_parser.parse_args(args) InfectionMonkey._log_arguments(opts) @@ -147,9 +153,9 @@ 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._opts.servers)}") + logger.debug(f"Trying to wake up with servers: {', '.join(self._server_strings)}") server_clients = find_available_island_apis( - self._opts.servers, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) + self._server_strings, HTTPIslandAPIClientFactory(self._agent_event_serializer_registry) ) server, island_api_client = self._select_server(server_clients) @@ -158,13 +164,13 @@ 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._opts.servers}" + f"Failed to connect to the island via any known servers: {self._server_strings}" ) # 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 and server_clients[s]) + servers_to_close = (s for s in self._server_strings if s != server and server_clients[s]) send_remove_from_waitlist_control_message_to_relays(servers_to_close) return server, island_api_client @@ -184,7 +190,7 @@ class InfectionMonkey: def _select_server( self, server_clients: Mapping[str, Optional[IIslandAPIClient]] ) -> Tuple[Optional[str], Optional[IIslandAPIClient]]: - for server in self._opts.servers: + for server in self._server_strings: if server_clients[server]: return server, server_clients[server] @@ -192,7 +198,7 @@ class InfectionMonkey: @staticmethod def _log_arguments(args): - arg_string = " ".join([f"{key}: {value}" for key, value in vars(args).items()]) + arg_string = ", ".join([f"{key}: {value}" for key, value in vars(args).items()]) logger.info(f"Monkey started with arguments: {arg_string}") def start(self): @@ -255,7 +261,7 @@ class InfectionMonkey: def _build_server_list(self, relay_port: int): relay_servers = [f"{ip}:{relay_port}" for ip in get_my_ip_addresses()] - return self._opts.servers + relay_servers + return self._server_strings + relay_servers def _build_master(self, relay_port: int): servers = self._build_server_list(relay_port) From 8b8ef79e0ad154759c5eddb64a212b9fcf814f9f Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 12:32:00 +0000 Subject: [PATCH 3/3] Common: Move SocketAddress function into the class Moved socketaddress_from_string into SocketAddress.from_string --- monkey/common/types.py | 36 +++++++++++++++++++ monkey/common/types/__init__.py | 1 - monkey/common/types/types.py | 19 ---------- monkey/common/types/utils/__init__.py | 1 - monkey/common/types/utils/socket_address.py | 18 ---------- monkey/infection_monkey/monkey.py | 4 +-- .../{types/utils => }/test_socket_address.py | 9 +++-- 7 files changed, 42 insertions(+), 46 deletions(-) create mode 100644 monkey/common/types.py delete mode 100644 monkey/common/types/__init__.py delete mode 100644 monkey/common/types/types.py delete mode 100644 monkey/common/types/utils/__init__.py delete mode 100644 monkey/common/types/utils/socket_address.py rename monkey/tests/unit_tests/common/{types/utils => }/test_socket_address.py (66%) diff --git a/monkey/common/types.py b/monkey/common/types.py new file mode 100644 index 000000000..5f86d5060 --- /dev/null +++ b/monkey/common/types.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from ipaddress import IPv4Address +from uuid import UUID + +from pydantic import PositiveInt, conint +from typing_extensions import TypeAlias + +from common.base_models import InfectionMonkeyBaseModel +from common.network.network_utils import address_to_ip_port + +AgentID: TypeAlias = UUID +HardwareID: TypeAlias = PositiveInt +MachineID: TypeAlias = PositiveInt + + +class SocketAddress(InfectionMonkeyBaseModel): + ip: IPv4Address + port: conint(ge=1, le=65535) # type: ignore[valid-type] + + @classmethod + def from_string(cls, address_str: str) -> SocketAddress: + """ + Parse a SocketAddress object from a string + + :param address_str: A string of ip:port + :raises ValueError: If the string is not a valid ip:port + :return: SocketAddress with the IP and port + """ + ip, port = address_to_ip_port(address_str) + if port is None: + raise ValueError("SocketAddress requires a port") + return SocketAddress(ip=IPv4Address(ip), port=int(port)) + + def __str__(self): + return f"{self.ip}:{self.port}" diff --git a/monkey/common/types/__init__.py b/monkey/common/types/__init__.py deleted file mode 100644 index 97bfe1c80..000000000 --- a/monkey/common/types/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .types import AgentID, HardwareID, MachineID, SocketAddress diff --git a/monkey/common/types/types.py b/monkey/common/types/types.py deleted file mode 100644 index 51353d293..000000000 --- a/monkey/common/types/types.py +++ /dev/null @@ -1,19 +0,0 @@ -from ipaddress import IPv4Address -from uuid import UUID - -from pydantic import PositiveInt, conint -from typing_extensions import TypeAlias - -from common.base_models import InfectionMonkeyBaseModel - -AgentID: TypeAlias = UUID -HardwareID: TypeAlias = PositiveInt -MachineID: TypeAlias = PositiveInt - - -class SocketAddress(InfectionMonkeyBaseModel): - ip: IPv4Address - port: conint(ge=1, le=65535) # type: ignore[valid-type] - - def __str__(self): - return f"{self.ip}:{self.port}" diff --git a/monkey/common/types/utils/__init__.py b/monkey/common/types/utils/__init__.py deleted file mode 100644 index 7d7440ee3..000000000 --- a/monkey/common/types/utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .socket_address import socketaddress_from_string diff --git a/monkey/common/types/utils/socket_address.py b/monkey/common/types/utils/socket_address.py deleted file mode 100644 index 77418d4be..000000000 --- a/monkey/common/types/utils/socket_address.py +++ /dev/null @@ -1,18 +0,0 @@ -from ipaddress import IPv4Address - -from common.network.network_utils import address_to_ip_port -from common.types import SocketAddress - - -def socketaddress_from_string(address_str: str) -> SocketAddress: - """ - Parse a SocketAddress object from a string - - :param address_str: A string of ip:port - :raises ValueError: If the string is not a valid ip:port - :return: SocketAddress with the IP and port - """ - ip, port = address_to_ip_port(address_str) - if port is None: - raise ValueError("SocketAddress requires a port") - return SocketAddress(ip=IPv4Address(ip), port=int(port)) diff --git a/monkey/infection_monkey/monkey.py b/monkey/infection_monkey/monkey.py index 1895211ea..73cd4e50a 100644 --- a/monkey/infection_monkey/monkey.py +++ b/monkey/infection_monkey/monkey.py @@ -21,7 +21,7 @@ from common.network.network_utils import ( get_my_ip_addresses, get_network_interfaces, ) -from common.types.utils import socketaddress_from_string +from common.types import SocketAddress from common.utils.argparse_types import positive_int from common.utils.attack_utils import ScanStatus, UsageEnum from common.version import get_version @@ -143,7 +143,7 @@ class InfectionMonkey: arg_parser.add_argument( "-s", "--servers", - type=lambda arg: [socketaddress_from_string(s) for s in arg.strip().split(",")], + type=lambda arg: [SocketAddress.from_string(s) for s in arg.strip().split(",")], ) arg_parser.add_argument("-d", "--depth", type=positive_int, default=0) opts = arg_parser.parse_args(args) diff --git a/monkey/tests/unit_tests/common/types/utils/test_socket_address.py b/monkey/tests/unit_tests/common/test_socket_address.py similarity index 66% rename from monkey/tests/unit_tests/common/types/utils/test_socket_address.py rename to monkey/tests/unit_tests/common/test_socket_address.py index a73fe8c6f..6c3b2d448 100644 --- a/monkey/tests/unit_tests/common/types/utils/test_socket_address.py +++ b/monkey/tests/unit_tests/common/test_socket_address.py @@ -1,7 +1,6 @@ import pytest from common.types import SocketAddress -from common.types.utils import socketaddress_from_string GOOD_IP = "192.168.1.1" BAD_IP = "192.168.1.999" @@ -9,10 +8,10 @@ GOOD_PORT = 1234 BAD_PORT = 99999 -def test_socketaddress_from_string(): +def test_socket_address__from_string(): expected = SocketAddress(ip=GOOD_IP, port=GOOD_PORT) - address = socketaddress_from_string(f"{GOOD_IP}:{GOOD_PORT}") + address = SocketAddress.from_string(f"{GOOD_IP}:{GOOD_PORT}") assert address == expected @@ -30,6 +29,6 @@ def test_socketaddress_from_string(): f"{GOOD_IP}:{BAD_PORT}", ], ) -def test_socketaddress_from_string__raises(bad_address: str): +def test_socket_address__from_string_raises(bad_address: str): with pytest.raises(ValueError): - socketaddress_from_string(bad_address) + SocketAddress.from_string(bad_address)