From 8b8ef79e0ad154759c5eddb64a212b9fcf814f9f Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
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)