Update bot/integrations/diag_extra.py

bot/integrations/diag_extra.py

@@ -1,52 +1,37 @@
 # PATH: bot/integrations/diag_extra.py
-from __future__ import annotations
-
 import socket
 from urllib.parse import urlparse
 
+def _host_from_url(url: str) -> str:
+    try:
+        u = urlparse(url)
+        if u.hostname:
+            return u.hostname
+    except Exception:
+        pass
+    # fallback: if someone passed "domain.com"
+    url = (url or "").strip()
+    url = url.replace("https://", "").replace("http://", "")
+    return url.split("/")[0].strip()
 
-def _to_host(url_or_host: str) -> str:
-    s = (url_or_host or "").strip()
-    if not s:
-        return ""
-    # If it's a URL, extract hostname; else treat as hostname
-    if "://" in s:
-        try:
-            u = urlparse(s)
-            return (u.hostname or "").strip()
-        except Exception:
-            return ""
-    # remove port if present
-    if ":" in s:
-        return s.split(":", 1)[0].strip()
-    return s
-
-
-def dns_check(url_or_host: str, max_ips: int = 5) -> str:
+def dns_check(url_or_host: str, port: int = 443) -> str:
     """
-    Returns a short string:
-      - OK:host -> ip1,ip2
-      - ERR:<Reason>
-    Works in HF container (best-effort).
+    Returns:
+      - OK:<host> -> ip1, ip2, ...
+      - ERR:<ExceptionName>:<message>
     """
-    host = _to_host(url_or_host)
+    host = _host_from_url(url_or_host)
     if not host:
-        return "ERR:empty_host"
+        return "ERR:ValueError:empty_host"
 
     try:
-        # getaddrinfo may return duplicates; we de-dupe
-        infos = socket.getaddrinfo(host, 443, type=socket.SOCK_STREAM)
+        infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
         ips = []
-        for info in infos:
-            ip = info[4][0]
-            if ip not in ips:
-                ips.append(ip)
-            if len(ips) >= max_ips:
-                break
-        if not ips:
-            return f"ERR:no_ip:{host}"
-        return f"OK:{host} -> {', '.join(ips)}"
-    except socket.gaierror as e:
-        return f"ERR:gaierror:{e}"
+        for it in infos:
+            addr = it[4][0]
+            if addr not in ips:
+                ips.append(addr)
+        ips = ips[:8]  # keep short
+        return f"OK:{host} -> " + ", ".join(ips)
     except Exception as e:
         return f"ERR:{type(e).__name__}:{e}"