From 4982999b999c19ecb1108c35ce7342fb14a1a088 Mon Sep 17 00:00:00 2001
From: Kekoa Kaaikala <kekoa.kaaikala@gmail.com>
Date: Fri, 23 Sep 2022 17:54:49 +0000
Subject: [PATCH] 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)