Update app.py

app.py

@@ -1,6 +1,5 @@
 import time
 import re
-import socket
 import ipaddress
 from urllib.parse import urlparse
 
@@ -8,6 +7,9 @@ import httpx
 import dns.resolver
 import gradio as gr
 
+# -----------------------
+# Safety: SSRF protection
+# -----------------------
 PRIVATE_NETS = [
     ipaddress.ip_network("0.0.0.0/8"),
     ipaddress.ip_network("10.0.0.0/8"),
@@ -15,44 +17,53 @@ PRIVATE_NETS = [
     ipaddress.ip_network("169.254.0.0/16"),
     ipaddress.ip_network("172.16.0.0/12"),
     ipaddress.ip_network("192.168.0.0/16"),
-    ipaddress.ip_network("224.0.0.0/4"),
-    ipaddress.ip_network("240.0.0.0/4"),
+    ipaddress.ip_network("224.0.0.0/4"),   # multicast
+    ipaddress.ip_network("240.0.0.0/4"),   # reserved
     ipaddress.ip_network("::1/128"),
-    ipaddress.ip_network("fc00::/7"),
-    ipaddress.ip_network("fe80::/10"),
+    ipaddress.ip_network("fc00::/7"),      # unique local
+    ipaddress.ip_network("fe80::/10"),     # link-local
 ]
 
-DOMAIN_RE = re.compile(r"^(?=.{1,253}$)(?!-)[A-Za-z0-9-]{1,63}(?<!-)(\.(?!-)[A-Za-z0-9-]{1,63}(?<!-))*$")
+# Basic domain sanity (not perfect IDN validation, but enough for UI)
+DOMAIN_RE = re.compile(
+    r"^(?=.{1,253}$)(?!-)[A-Za-z0-9-]{1,63}(?<!-)(\.(?!-)[A-Za-z0-9-]{1,63}(?<!-))*$"
+)
 
-# ✅ Reuse one client (less event-loop weirdness + faster)
+# Reuse one client (faster + avoids repeated loop churn)
 CLIENT = httpx.Client(
     follow_redirects=True,
     timeout=8.0,
     headers={"User-Agent": "HF-Connectivity-Checker/1.0"},
 )
 
-def _is_private_ip(ip_str: str) -> bool:
+def is_private_ip(ip_str: str) -> bool:
     try:
         ip = ipaddress.ip_address(ip_str)
         return any(ip in net for net in PRIVATE_NETS)
     except Exception:
         return True
 
-def _parse_target(target: str):
+def parse_target(target: str):
+    """
+    Returns: (kind, raw, host)
+      kind: url | domain | ip | unknown | empty
+    """
     t = (target or "").strip()
     if not t:
-        return ("", "", "")
+        return ("empty", "", "")
 
     if t.startswith("http://") or t.startswith("https://"):
         u = urlparse(t)
         return ("url", t, u.hostname or "")
 
+    # IP?
     try:
         ipaddress.ip_address(t)
         return ("ip", t, t)
     except Exception:
         pass
 
+    # Domain?
     d = t.rstrip(".")
     if DOMAIN_RE.match(d):
         return ("domain", d, d)
@@ -61,6 +72,7 @@ def _parse_target(target: str):
 
 def resolve_dns(host: str):
     out = {"A": [], "AAAA": [], "CNAME": []}
+
     if not host:
         return out
 
@@ -86,9 +98,11 @@ def try_http(url: str):
     try:
         start = time.time()
         r = CLIENT.head(url)
+        # Some sites block HEAD; fallback to tiny GET
         if r.status_code in (405, 403) or r.status_code >= 500:
             r = CLIENT.get(url, headers={"Range": "bytes=0-1024"})
         elapsed = int((time.time() - start) * 1000)
+
         info.update({
             "ok": True,
             "status_code": r.status_code,
@@ -111,9 +125,11 @@ def try_http(url: str):
     return info
 
 def classify(dns_info, http_info):
+    # DNS fail
     if dns_info.get("A_error") and dns_info.get("AAAA_error") and not dns_info.get("CNAME"):
         return "DNS_FAIL (HF can't resolve)"
 
+    # HTTP ok?
     ok = [x for x in http_info if x.get("ok")]
     if ok:
         code = ok[0].get("status_code", 0)
@@ -131,54 +147,55 @@ def classify(dns_info, http_info):
     return f"NOT_REACHABLE ({errs or 'unknown'})"
 
 def check_one(target: str):
-    kind, raw, host = _parse_target(target)
+    kind, raw, host = parse_target(target)
 
-    if kind == "":
+    if kind == "empty":
         return {"error": "Enter a domain / IP / URL"}
     if kind == "unknown" or not host:
         return {"error": "Invalid input"}
 
-    # IP safety
+    # IP checks
     try:
         ipaddress.ip_address(host)
-        if _is_private_ip(host):
-            return {"error": "Blocked: private/reserved IP not allowed."}
+        if is_private_ip(host):
+            return {"error": "Blocked: private/reserved IP not allowed (SSRF protection)."}
         dns_info = {"A": [host], "AAAA": [], "CNAME": []}
         host_for_http = host
     except Exception:
         dns_info = resolve_dns(host)
         ips = (dns_info.get("A") or []) + (dns_info.get("AAAA") or [])
         for ip in ips:
-            if _is_private_ip(ip):
+            if is_private_ip(ip):
                 return {"error": "Blocked: resolves to private/reserved IP (SSRF protection)."}
         host_for_http = host
 
+    # Only 2 tries to keep it safe & fast
     urls = []
     if kind == "url":
         urls.append(raw)
     urls.append(f"https://{host_for_http}")
     urls.append(f"http://{host_for_http}")
 
-    http_results = [try_http(urls[0]), try_http(urls[1])]  # only two tries
+    http_results = [try_http(urls[0]), try_http(urls[1])]
     status = classify(dns_info, http_results)
 
     return {
-        "input": target.strip(),
+        "input": (target or "").strip(),
         "host": host_for_http,
         "dns": dns_info,
         "http": http_results,
         "status": status,
-        "note": "Checked from Hugging Face Space network."
+        "note": "Checked from Hugging Face Space network (egress).",
     }
 
-def bulk_check(domain: str, subdomains_text: str):
-    base = (domain or "").strip().rstrip(".")
+def bulk_check(base_domain: str, subdomains_text: str):
+    base = (base_domain or "").strip().rstrip(".")
     if not base:
         return []
 
     lines = [x.strip() for x in (subdomains_text or "").splitlines() if x.strip()]
     targets = []
-    for s in lines[:200]:
+    for s in lines[:200]:  # limit
         targets.append(s if "." in s else f"{s}.{base}")
 
     rows = []
@@ -186,21 +203,21 @@ def bulk_check(domain: str, subdomains_text: str):
         r = check_one(t)
         dns = r.get("dns", {}) if isinstance(r, dict) else {}
         http = r.get("http", [{}]) if isinstance(r, dict) else [{}]
-        rows.append({
-            "target": t,
-            "status": r.get("status") or r.get("error", "error"),
-            "A": ",".join(dns.get("A", [])),
-            "AAAA": ",".join(dns.get("AAAA", [])),
-            "http_code": http[0].get("status_code", ""),
-        })
+        rows.append([
+            t,
+            r.get("status") or r.get("error", "error"),
+            ",".join(dns.get("A", [])),
+            ",".join(dns.get("AAAA", [])),
+            str(http[0].get("status_code", "")),
+        ])
     return rows
 
-with gr.Blocks(title="HF Domain/IP Connectivity Checker") as demo:
+with gr.Blocks(title="HF Domain IP Checker") as demo:
     gr.Markdown(
         "## HF Domain/IP Connectivity Checker\n"
-        "Checks DNS + HTTP reachability **from this Hugging Face Space**.\n"
-        "- Subdomains must be **your provided list** (no brute-force scanning).\n"
-        "- Private/reserved IPs blocked (SSRF protection)."
+        "DNS + HTTP reachability **from this Hugging Face Space**.\n"
+        "- Subdomains are checked only from your provided list.\n"
+        "- Private/reserved IPs are blocked (SSRF protection)."
     )
 
     with gr.Tab("Single Check"):
@@ -217,10 +234,9 @@ with gr.Blocks(title="HF Domain/IP Connectivity Checker") as demo:
             headers=["target", "status", "A", "AAAA", "http_code"],
             datatype=["str", "str", "str", "str", "str"],
             row_count=5,
-            column_count=(5, "fixed"),  # ✅ new param
             label="Results",
         )
         btn2.click(bulk_check, inputs=[base, subs], outputs=table)
 
-# ✅ SSR off to avoid Python 3.13 loop cleanup errors
+# SSR off for stability
 demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)