#!/usr/bin/env python3
"""ssrf_prescan.py — deterministic URL pre-scan for SSRF-risk surface.

ARION service artifact. Given a URL, reports whether its host is a literal or
resolves to private / reserved / non-global address space — including the
bypass classes that beat naive `isPrivateAddress` guards:

  * non-canonical IPv4 spellings: single int (2130706433), hex (0x7f000001),
    octal (0177.0.0.1), short forms (127.1)  [inet_aton semantics]
  * IPv6-embedded IPv4: ::ffff:0:0/96 (mapped), 64:ff9b::/96 + 64:ff9b:1::/48
    (NAT64/DNS64 — the vector behind agenticjobs bug #7), 2002::/16 (6to4),
    2001::/32 (Teredo)
  * unspecified/compat :: and ::1
  * per-hop redirect re-check (--follow): each Location hop is re-scanned,
    matching the redirect-recheck fix

Honest scope: this is a PRE-SCAN, not a proof of exploitability and not a
full audit. DNS answers can differ at fetch time (rebinding/TOCTOU) — a PASS
verdict means "nothing flagged at scan time", not "safe forever".

Usage:
  ssrf_prescan.py <url> [--follow] [--no-resolve] [--max-hops N] [--json]
  ssrf_prescan.py --selftest
Stdlib only. Exit 0 = no flag, 1 = at least one flag, 2 = usage/scan error.
"""
import ipaddress
import json
import socket
import ssl
import sys
import urllib.error
import urllib.parse
import urllib.request

V4_NETWORKS = [
    ("0.0.0.0/8", "this-host/broadcast-src"),
    ("10.0.0.0/8", "RFC1918 private"),
    ("100.64.0.0/10", "CGNAT shared"),
    ("127.0.0.0/8", "loopback"),
    ("169.254.0.0/16", "link-local (cloud metadata 169.254.169.254)"),
    ("172.16.0.0/12", "RFC1918 private"),
    ("192.0.0.0/24", "IETF protocol assignments"),
    ("192.0.2.0/24", "TEST-NET-1 doc"),
    ("192.168.0.0/16", "RFC1918 private"),
    ("198.18.0.0/15", "benchmarking"),
    ("198.51.100.0/24", "TEST-NET-2 doc"),
    ("203.0.113.0/24", "TEST-NET-3 doc"),
    ("224.0.0.0/4", "multicast"),
    ("240.0.0.0/4", "reserved"),
]
V6_NETWORKS = [
    ("::/128", "unspecified"),
    ("::1/128", "loopback"),
    ("64:ff9b:1::/48", "NAT64 local-use (embedded v4 judged)"),
    ("100::/64", "discard-only"),
    ("2001:db8::/32", "documentation"),
    ("fc00::/7", "ULA private"),
    ("fe80::/10", "link-local"),
    ("ff00::/8", "multicast"),
]
V4_NETS = [(ipaddress.ip_network(n), w) for n, w in V4_NETWORKS]
V6_NETS = [(ipaddress.ip_network(n), w) for n, w in V6_NETWORKS]

NAT64 = ipaddress.ip_network("64:ff9b::/96")
V4MAPPED = ipaddress.ip_network("::ffff:0:0/96")
SIXTO4 = ipaddress.ip_network("2002::/16")
TEREDO = ipaddress.ip_network("2001::/32")


def v4_reason(ip):
    for net, why in V4_NETS:
        if ip in net:
            return f"{ip} in {net.with_prefixlen} ({why})"
    return None


def embedded_v4(ip6):
    """Return (source, ipv4) for the classes that smuggle v4 inside v6."""
    raw = ip6.packed
    if ip6 in V4MAPPED:
        return "v4-mapped ::ffff:0:0/96", ipaddress.ip_address(raw[12:])
    if ip6 in NAT64:
        return "NAT64 64:ff9b::/96", ipaddress.ip_address(raw[12:])
    if ip6 in SIXTO4:
        return "6to4 2002::/16", ipaddress.ip_address(raw[2:6])
    if ip6 in TEREDO:
        return ("Teredo 2001::/32 (v4 xor-encoded)",
                ipaddress.ip_address(bytes(b ^ 0xFF for b in raw[12:16])))
    return None, None


def ip_reason(ip):
    """Flag reason string, or None if the address is public-global-looking."""
    if ip.version == 4:
        return v4_reason(ip)
    src, v4 = embedded_v4(ip)
    if v4 is not None:
        r = v4_reason(v4)
        if r:
            return f"{ip} {src} embeds {r}"
        return None  # embedded v4 is global — not itself a flag
    for net, why in V6_NETS:
        if ip in net:
            return f"{ip} in {net.with_prefixlen} ({why})"
    return None


def aton(host):
    """inet_aton-style parse: 1-4 dot-separated decimal/hex/octal parts."""
    parts = host.split(".")
    if len(parts) > 4 or any(p == "" for p in parts):
        return None
    vals = []
    for p in parts:
        try:
            if p.lower().startswith("0x"):
                v = int(p, 16)
            elif len(p) > 1 and p.startswith("0"):
                v = int(p, 8)
            else:
                v = int(p, 10)
        except ValueError:
            return None
        vals.append(v)
    if vals[-1] >= 256 ** (5 - len(vals)) or any(v > 255 for v in vals[:-1]):
        return None
    n = vals[-1]
    for i, v in enumerate(vals[:-1]):
        n += v << (8 * (3 - i))
    return ipaddress.ip_address(n)


def classify_host(host):
    """Return (kind, ip|None). kind: 'v4','v6','dns','bad'."""
    if host.startswith("[") and host.endswith("]"):
        try:
            return "v6", ipaddress.ip_address(host[1:-1])
        except ValueError:
            return "bad", None
    try:
        return "v4", ipaddress.ip_address(host)
    except ValueError:
        pass
    ip = aton(host)
    if ip is not None:
        return "v4", ip
    return "dns", None


def resolve(host):
    fams = set()
    try:
        for _af, _t, _p, _c, sa in socket.getaddrinfo(
                host, None, proto=socket.IPPROTO_TCP):
            fams.add(sa[0])
    except socket.gaierror as e:
        return None, f"DNS resolution failed: {e}"
    out = []
    for a in sorted(fams):
        try:
            out.append(ipaddress.ip_address(a))
        except ValueError:
            pass
    return out, None


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *a, **k):
        return None


class _HeadRequest(urllib.request.Request):
    def get_method(self):
        return "HEAD"


def fetch_redirects(url, max_hops):
    """Walk redirects WITHOUT auto-following; return list of hop URLs."""
    ctx = ssl.create_default_context()
    opener = urllib.request.build_opener(
        _NoRedirect, urllib.request.HTTPSHandler(context=ctx))
    hops = [url]
    for _ in range(max_hops):
        cur = hops[-1]
        try:
            opener.open(_HeadRequest(cur), timeout=10)
        except urllib.error.HTTPError as e:
            if e.is_redirect() if hasattr(e, "is_redirect") else \
                    e.code in (301, 302, 303, 307, 308):
                loc = e.headers.get("Location")
                if not loc:
                    break
                hops.append(urllib.parse.urljoin(cur, loc))
                continue
            hops.append(f"(http {e.code})")
            break
        except Exception as e:
            hops.append(f"(fetch error: {e})")
            break
        else:
            break
    return hops


def scan(url, do_resolve=True):
    """Scan one URL. Returns dict with verdict + findings."""
    row = {"url": url, "flags": [], "notes": [], "ips": []}
    try:
        p = urllib.parse.urlsplit(url)
    except ValueError as e:
        row["flags"].append(f"unparseable URL: {e}")
        row["verdict"] = "FLAG"
        return row
    if p.scheme not in ("http", "https"):
        row["flags"].append(
            f"scheme '{p.scheme or '(none)'}' outside http/https allowlist")
    if p.username or p.password:
        row["notes"].append("userinfo present (possible authority confusion)")
    host = p.hostname or ""
    if not host:
        row["flags"].append("no host")
        row["verdict"] = "FLAG"
        return row
    port = p.port
    if port is not None:
        row["notes"].append(f"explicit port {port}")
    kind, ip = classify_host(host)
    row["host_kind"] = kind
    ips = []
    if ip is not None:
        ips = [ip]
        if kind == "v4" and host != str(ip):
            row["notes"].append(
                f"non-canonical IPv4 spelling '{host}' = {ip}")
    elif kind == "dns":
        if do_resolve:
            ips, err = resolve(host.rstrip("."))
            if err:
                row["notes"].append(err)
        else:
            row["notes"].append("DNS name (resolution skipped)")
    row["ips"] = [str(i) for i in ips]
    for i in ips:
        r = ip_reason(i)
        if r:
            row["flags"].append(f"{host} -> {r}")
    row["verdict"] = "FLAG" if row["flags"] else "PASS"
    return row


def report(results):
    out = ["# SSRF pre-scan report",
           "# deterministic pre-scan only — DNS may change at fetch time "
           "(rebind/TOCTOU); PASS != guaranteed safe", ""]
    worst = "PASS"
    for r in results:
        mark = "FLAG" if r["verdict"] == "FLAG" else "pass"
        if r["verdict"] == "FLAG":
            worst = "FLAG"
        out.append(f"## [{mark}] {r['url']}")
        for i in r.get("ips", []):
            out.append(f"  resolved/literal: {i}")
        for n in r.get("notes", []):
            out.append(f"  note: {n}")
        for f in r.get("flags", []):
            out.append(f"  FLAG: {f}")
        out.append("")
    out.append(f"overall: {worst}")
    return "\n".join(out)


SELFTEST = [
    ("http://127.0.0.1/", "FLAG"),
    ("http://2130706433/", "FLAG"),          # 127.0.0.1 as dword
    ("http://0x7f000001/", "FLAG"),          # hex
    ("http://0177.0.0.1/", "FLAG"),          # octal
    ("http://127.1/", "FLAG"),               # short form
    ("http://0/", "FLAG"),                   # 0.0.0.0
    ("http://10.0.0.1/", "FLAG"),
    ("http://169.254.169.254/", "FLAG"),     # cloud metadata
    ("http://[::1]/", "FLAG"),
    ("http://[::]/", "FLAG"),
    ("http://[::ffff:127.0.0.1]/", "FLAG"),  # mapped
    ("http://[64:ff9b::7f00:1]/", "FLAG"),   # NAT64 embeds 127.0.0.1
    ("http://[64:ff9b::a00:1]/", "FLAG"),    # NAT64 embeds 10.0.0.1
    ("http://[2002:7f00:1::]/", "FLAG"),     # 6to4 embeds 127.0.0.1
    ("http://[2001:0:4136:e378:8000:63bf:3fff:fdd2]/", "FLAG"),  # Teredo
    ("http://1.1.1.1/", "PASS"),
    ("http://[2606:4700:4700::1111]/", "PASS"),
    ("http://[64:ff9b::808:808]/", "PASS"),  # NAT64 embeds 8.8.8.8 (global)
    ("http://[2002:808:808::]/", "PASS"),    # 6to4 embeds 8.8.8.8
    ("file:///etc/passwd", "FLAG"),          # scheme
    ("gopher://x/", "FLAG"),
]


def selftest():
    fails = 0
    for url, want in SELFTEST:
        got = scan(url, do_resolve=False)["verdict"]
        ok = got == want
        fails += not ok
        print(f"{'ok ' if ok else 'FAIL'} {want} {url} -> {got}")
    print(f"{len(SELFTEST) - fails}/{len(SELFTEST)} pass")
    return 1 if fails else 0


def main(argv):
    args = [a for a in argv if not a.startswith("--")]
    flags = {a for a in argv if a.startswith("--")}
    if "--selftest" in flags:
        return selftest()
    if not args:
        print(__doc__.splitlines()[0])
        print("usage: ssrf_prescan.py <url> [--follow] [--max-hops N] "
              "[--no-resolve] [--json]")
        return 2
    max_hops = 5
    for f in list(flags):
        if f.startswith("--max-hops"):
            try:
                max_hops = int(f.split("=", 1)[1])
            except (IndexError, ValueError):
                pass
    urls = [args[0]]
    if "--follow" in flags:
        try:
            urls = fetch_redirects(args[0], max_hops)
        except Exception as e:
            print(f"redirect walk failed: {e}", file=sys.stderr)
    results = [scan(u, "--no-resolve" not in flags) for u in urls]
    if "--json" in flags:
        print(json.dumps(results, indent=1))
    else:
        print(report(results))
    return 1 if any(r["verdict"] == "FLAG" for r in results) else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))