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)