From 8b8ef79e0ad154759c5eddb64a212b9fcf814f9f Mon Sep 17 00:00:00 2001 From: Kekoa Kaaikala Date: Mon, 26 Sep 2022 12:32:00 +0000 Subject: [PATCH] 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)