diff --git a/aikido_zen/storage/hostnames.py b/aikido_zen/storage/hostnames.py index 6b3d86a02..22b14dac1 100644 --- a/aikido_zen/storage/hostnames.py +++ b/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 diff --git a/aikido_zen/storage/hostnames_test.py b/aikido_zen/storage/hostnames_test.py index b7a6bea35..92218e01e 100644 --- a/aikido_zen/storage/hostnames_test.py +++ b/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)