Common: Move SocketAddress function into the class

Moved socketaddress_from_string into SocketAddress.from_string

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}"

monkey/common/types/__init__.py

@@ -1 +0,0 @@
-from .types import AgentID, HardwareID, MachineID, SocketAddress

monkey/common/types/types.py

@@ -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}"

monkey/common/types/utils/__init__.py

@@ -1 +0,0 @@
-from .socket_address import socketaddress_from_string

monkey/common/types/utils/socket_address.py

@@ -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))

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)

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)