Fix reporting of port as string for outgoing domains

aikido_zen/storage/hostnames.py

@@ -1,5 +1,7 @@
 """Export Hostnames class"""
 
+import socket
+
 
 class Hostnames:
     """Stores hostnames"""
@@ -10,9 +12,10 @@ def __init__(self, max_entries=200):
 
     def add(self, hostname, port, hits=1):
         """Add a hostname and port to the map"""
-        key = get_key(hostname, port)
+        normalized_port = normalize_port(port)
+        key = get_key(hostname, normalized_port)
         if not self.map.get(key):
-            self.map[key] = {"hostname": hostname, "port": port, "hits": 0}
+            self.map[key] = {"hostname": hostname, "port": normalized_port, "hits": 0}
         if len(self.map) > self.max_entries:
             # Remove the first added hostname
             first_added = next(iter(self.map))
@@ -31,3 +34,21 @@ def clear(self):
 def get_key(hostname, port):
     """Returns a string key"""
     return f"{hostname}:{port}"
+
+
+def normalize_port(port):
+    """
+    Ensures port is an int (or None), never a str.
+    `socket.getaddrinfo` accepts a service name (e.g. "http") or a numeric
+    string as the port, so it needs to be resolved/parsed to a number here.
+    """
+    if port is None or isinstance(port, int):
+        return port
+    try:
+        return int(port)
+    except (TypeError, ValueError):
+        pass
+    try:
+        return socket.getservbyname(port)
+    except OSError:
+        return None

aikido_zen/storage/hostnames_test.py

@@ -99,6 +99,35 @@ def test_add_none_port():
     assert hostnames.map[key]["hits"] == 1
 
 
+def test_add_numeric_string_port():
+    """Test that a numeric string port is stored as an int."""
+    hostnames = Hostnames(max_entries=3)
+    hostnames.add("example.com", "80")
+    key = "example.com:80"
+    assert key in hostnames.map
+    assert hostnames.map[key]["port"] == 80
+    assert isinstance(hostnames.map[key]["port"], int)
+
+
+def test_add_service_name_port():
+    """Test that a service name port (e.g. 'https') is resolved to its number."""
+    hostnames = Hostnames(max_entries=3)
+    hostnames.add("example.com", "https")
+    key = "example.com:443"
+    assert key in hostnames.map
+    assert hostnames.map[key]["port"] == 443
+    assert isinstance(hostnames.map[key]["port"], int)
+
+
+def test_add_unknown_service_name_port():
+    """Test that an unresolvable service name port falls back to None."""
+    hostnames = Hostnames(max_entries=3)
+    hostnames.add("example.com", "not-a-real-service")
+    key = "example.com:None"
+    assert key in hostnames.map
+    assert hostnames.map[key]["port"] is None
+
+
 def test_exceed_max_entries_with_multiple_ports():
     """Test exceeding max entries with multiple ports."""
     hostnames = Hostnames(max_entries=3)