#!/usr/bin/env python3
"""The probe harness. It runs INSIDE a `docker run --network none` container (see run-arm.sh).

1. Starts three loopback listeners that log everything and answer as little as possible:
     DNS  127.0.0.1:53  (UDP+TCP): logs (time, name, type); every A query gets 127.0.0.2, all else NOERROR/empty.
     TLS  127.0.0.2:443: logs the connection, parses the ClientHello (SNI, ALPN), closes WITHOUT answering.
     HTTP 127.0.0.2:80:  logs the request line and Host header, closes.
2. Self-test (a failure voids the run): public IPs unreachable, our DNS answers, our TLS listener sees SNI.
3. Starts the subject (subject.py) under strace (-f, %network) with a from-scratch environment.
4. After the subject exits (+ a short settle), lists every file under /tmp (which holds HOME), copies the
   ORT queue DB if one appeared, reads its rows, and writes everything to /probe/out/result.json.

Standard library only. Nothing here talks to anything but loopback.
"""
import argparse
import hashlib
import json
import os
import re
import shutil
import socket
import sqlite3
import ssl
import struct
import subprocess
import sys
import threading
import time

OUT = "/probe/out"
HOME = "/tmp/home"
QUEUE_DIR = HOME + "/.cache/Microsoft/DeveloperTools/.onnxruntime"
SINK_IP = "127.0.0.2"          # every name resolves here
SELFTEST_NAME = "probe-selftest.example"
COLLECTOR = "mobile.events.data.microsoft.com"
PUBLIC_IPS = ["1.1.1.1", "20.184.175.9", "20.184.175.13"]  # a resolver + two addresses the collector has used
CANARIES = ["probe_graph_canary", "probe_meta_key", "probe_meta_value_canary", "probe-kit"]
ORT_EVENTS = [  # event names in onnxruntime 1.30.0 posix/telemetry.cc (recon section 2.2)
    "ProcessInfo", "RegisterEpLibraryStart", "RegisterEpLibraryEnd", "RegisterEpLibraryWithLibPath",
    "SessionCreationStart", "SessionCreation", "SessionCreationEnd", "ModelLoadStart", "ModelLoadEnd",
    "EpDeviceUsage", "RuntimePerf", "SystemMetrics", "RuntimeError", "RuntimeInferenceError",
    "CompileModelStart", "CompileModelComplete", "EpAutoSelection",
]

EVENTS, LOCK = [], threading.Lock()


def log(kind, **kw):
    with LOCK:
        EVENTS.append({"t": time.time(), "kind": kind, **kw})


# ---------------------------------------------------------------- DNS
def dns_parse_question(q):
    i, labels = 12, []
    while q[i]:
        n = q[i]
        labels.append(q[i + 1:i + 1 + n].decode("ascii", "replace"))
        i += 1 + n
    i += 1
    qtype, _qclass = struct.unpack(">HH", q[i:i + 4])
    return ".".join(labels), qtype, q[12:i + 4]


def dns_reply(q, transport):
    name, qtype, question = dns_parse_question(q)
    log("dns", name=name, qtype={1: "A", 28: "AAAA"}.get(qtype, str(qtype)), transport=transport)
    flags = 0x8000 | ((q[2] & 0x01) << 8) | 0x0080          # response, copy RD, set RA, NOERROR
    answer = b""
    if qtype == 1:
        answer = b"\xc0\x0c" + struct.pack(">HHIH", 1, 1, 5, 4) + socket.inet_aton(SINK_IP)
    return q[:2] + struct.pack(">HHHHH", flags, 1, 1 if answer else 0, 0, 0) + question + answer


def dns_udp(sock):
    while True:
        q, peer = sock.recvfrom(4096)
        try:
            sock.sendto(dns_reply(q, "udp"), peer)
        except Exception as e:
            log("dns_error", error=repr(e))


def dns_tcp(srv):
    while True:
        conn, _ = srv.accept()
        with conn:
            try:
                n = struct.unpack(">H", conn.recv(2))[0]
                q = b""
                while len(q) < n:
                    q += conn.recv(n - len(q))
                r = dns_reply(q, "tcp")
                conn.sendall(struct.pack(">H", len(r)) + r)
            except Exception as e:
                log("dns_error", error=repr(e))


# ---------------------------------------------------------------- TLS / HTTP
def read_client_hello(conn):
    conn.settimeout(5)
    buf = b""
    try:
        while True:
            chunk = conn.recv(65536)
            if not chunk:
                break
            buf += chunk
            if len(buf) >= 5 and buf[0] == 0x16 and len(buf) >= 5 + int.from_bytes(buf[3:5], "big"):
                break
            if len(buf) >= 5 and buf[0] != 0x16:
                break
    except socket.timeout:
        pass
    return buf


def parse_client_hello(buf):
    """Pull SNI, ALPN and offered TLS versions out of a ClientHello. Returns {} if it is not one."""
    if len(buf) < 9 or buf[0] != 0x16 or buf[5] != 0x01:
        return {}
    body = buf[5:5 + int.from_bytes(buf[3:5], "big")]
    p = 4 + 2 + 32                                          # handshake header, client_version, random
    p += 1 + body[p]                                        # session id
    p += 2 + int.from_bytes(body[p:p + 2], "big")           # cipher suites
    p += 1 + body[p]                                        # compression methods
    end = p + 2 + int.from_bytes(body[p:p + 2], "big")
    p += 2
    info = {"sni": None, "alpn": [], "versions": []}
    while p + 4 <= end:
        et, el = struct.unpack(">HH", body[p:p + 4])
        ed = body[p + 4:p + 4 + el]
        p += 4 + el
        if et == 0 and len(ed) >= 5:                        # server_name
            info["sni"] = ed[5:5 + int.from_bytes(ed[3:5], "big")].decode("ascii", "replace")
        elif et == 16:                                      # ALPN
            j = 2
            while j < len(ed):
                info["alpn"].append(ed[j + 1:j + 1 + ed[j]].decode("ascii", "replace"))
                j += 1 + ed[j]
        elif et == 43 and ed:                               # supported_versions
            info["versions"] = [ed[k:k + 2].hex() for k in range(1, 1 + ed[0], 2)]
    return info


def tls_listener(srv):
    while True:
        conn, peer = srv.accept()
        with conn:
            t_accept = time.time()
            buf = read_client_hello(conn)
            log("tls", t_accept=t_accept, peer_port=peer[1], dst=conn.getsockname()[0],
                bytes=len(buf), **parse_client_hello(buf))
        # closing here, mid-handshake, is the point: no ServerHello, no session, nothing accepted


def http_listener(srv):
    while True:
        conn, peer = srv.accept()
        with conn:
            conn.settimeout(5)
            try:
                head = conn.recv(4096).decode("latin-1")
            except socket.timeout:
                head = ""
            host = next((l.split(":", 1)[1].strip() for l in head.split("\r\n") if l.lower().startswith("host:")), None)
            log("http", peer_port=peer[1], request_line=head.split("\r\n")[0], host=host)


def start_listeners():
    socks = {}
    u = socket.socket(socket.AF_INET, socket.SOCK_DGRAM); u.bind(("127.0.0.1", 53)); socks["dns_udp"] = u
    for key, addr, port in [("dns_tcp", "127.0.0.1", 53), ("tls", SINK_IP, 443), ("http", SINK_IP, 80)]:
        s = socket.socket(); s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind((addr, port)); s.listen(64); socks[key] = s
    for fn, key in [(dns_udp, "dns_udp"), (dns_tcp, "dns_tcp"), (tls_listener, "tls"), (http_listener, "http")]:
        threading.Thread(target=fn, args=(socks[key],), daemon=True).start()


# ---------------------------------------------------------------- self-test
def self_test():
    st = {"interfaces": sorted(os.listdir("/sys/class/net"))}
    with open("/proc/net/route") as f:
        st["ipv4_routes"] = len(f.read().splitlines()) - 1
    try:
        with open("/proc/net/ipv6_route") as f:
            st["ipv6_routes_not_lo"] = sum(1 for l in f if not l.rstrip().endswith(" lo"))
    except OSError:
        st["ipv6_routes_not_lo"] = "no ipv6"
    with open("/proc/sys/net/ipv4/ip_unprivileged_port_start") as f:
        st["ip_unprivileged_port_start"] = int(f.read())
    st["public_connects"] = {}
    for ip in PUBLIC_IPS:
        try:
            socket.create_connection((ip, 443), timeout=3).close()
            st["public_connects"][ip] = "CONNECTED"
        except OSError as e:
            st["public_connects"][ip] = f"errno {e.errno} {e.strerror}"
    try:
        st["selftest_resolves_to"] = sorted({a[4][0] for a in socket.getaddrinfo(SELFTEST_NAME, 443, socket.AF_INET)})
    except OSError as e:
        st["selftest_resolves_to"] = [repr(e)]
    try:
        with socket.create_connection((SELFTEST_NAME, 443), timeout=3) as raw:
            ssl.create_default_context().wrap_socket(raw, server_hostname=SELFTEST_NAME)
    except Exception as e:
        st["selftest_tls_client_error"] = repr(e)
    time.sleep(0.5)
    with LOCK:
        st["selftest_dns_logged"] = any(e["kind"] == "dns" and e["name"] == SELFTEST_NAME for e in EVENTS)
        st["selftest_sni_logged"] = any(e["kind"] == "tls" and e.get("sni") == SELFTEST_NAME for e in EVENTS)
    st["pass"] = (all(v.startswith("errno 101") for v in st["public_connects"].values())
                  and st["selftest_resolves_to"] == [SINK_IP]
                  and st["selftest_dns_logged"] and st["selftest_sni_logged"]
                  and st["interfaces"] == ["lo"] and st["ipv4_routes"] == 0
                  and st["ipv6_routes_not_lo"] in (0, "no ipv6"))
    return st


# ---------------------------------------------------------------- after the run
def list_files(root, t0):
    out = []
    for d, _dirs, files in os.walk(root):
        for n in files:
            p = os.path.join(d, n)
            try:
                s = os.lstat(p)
                out.append({"path": p, "size": s.st_size, "mtime_rel_s": round(s.st_mtime - t0, 3)})
            except OSError:
                pass
    return sorted(out, key=lambda x: x["path"])


def event_names(payload):
    """Known ORT event names present in a payload, longest match wins (SessionCreationStart > SessionCreation)."""
    text = payload.decode("latin-1")
    hits = {}
    for name in ORT_EVENTS:
        for m in re.finditer(re.escape(name), text):
            if len(name) > len(hits.get(m.start(), "")):
                hits[m.start()] = name
    return sorted(set(hits.values()))


def read_queue(t_launch, deviceid_raw, machine_id_raw, hostname):
    db = os.path.join(QUEUE_DIR, "onnxruntime.db")
    q = {"db_exists": os.path.exists(db)}
    if not q["db_exists"]:
        return q
    qdir = os.path.join(OUT, "queue")
    os.makedirs(qdir, exist_ok=True)
    for suffix in ("", "-wal", "-shm"):                     # copy first; never open the original
        if os.path.exists(db + suffix):
            shutil.copy2(db + suffix, os.path.join(qdir, "onnxruntime.db" + suffix))
    con = sqlite3.connect(os.path.join(qdir, "onnxruntime.db"))
    cols = [r[1] for r in con.execute("PRAGMA table_info(events)")]
    q["columns"] = cols
    dev_hash = ("c:" + hashlib.sha256(deviceid_raw.encode()).hexdigest().upper()) if deviceid_raw else None
    rows, strings = [], []
    for rec in con.execute("SELECT record_id, tenant_token, latency, persistence, timestamp, retry_count, "
                           "reserved_until, payload FROM events ORDER BY timestamp, record_id"):
        rid, tenant, latency, persistence, ts, retry, reserved, payload = rec
        payload = bytes(payload or b"")
        text = payload.decode("latin-1")
        rows.append({
            "tenant_prefix": (tenant or "")[:10], "latency": latency, "persistence": persistence,
            "timestamp_rel_s": round(ts / 1000 - t_launch, 3) if ts else None,
            "retry_count": retry, "payload_bytes": len(payload), "events": event_names(payload),
            "canaries": [c for c in CANARIES if c in text],
            "interpreter_path_in_payload": "/probe/work/venvs/" in text,
            "device_hash_c_sha256_deviceid": bool(dev_hash and dev_hash in text),
            "raw_deviceid": bool(deviceid_raw and deviceid_raw in text),
            "raw_machine_id": bool(machine_id_raw and machine_id_raw in text),
            "hostname": hostname in text,
        })
        strings.append(f"== record {rid}\n" + "\n".join(re.findall(r"[\x20-\x7e]{4,}", text)))
    q["settings"] = [list(r) for r in con.execute("SELECT name, value FROM settings")] if "settings" in {
        r[0] for r in con.execute("SELECT name FROM sqlite_master WHERE type='table'")} else None
    con.close()
    q["rows"] = rows
    with open(os.path.join(qdir, "strings.txt"), "w") as f:
        f.write("\n".join(strings) + "\n")
    return q


STRACE_RE = re.compile(r"^(\d+)\s+(\d+\.\d+)\s+(connect|sendto|sendmsg|sendmmsg)\((\d+),\s*(.*)$")


def parse_strace(path, t_launch):
    """Every network syscall that names an address, classified by family and destination."""
    calls = []
    if not os.path.exists(path):
        return None
    with open(path, errors="replace") as f:
        for line in f:
            m = STRACE_RE.match(line)
            if not m:
                continue
            pid, ts, call, _fd, rest = m.groups()
            c = {"t_rel_s": round(float(ts) - t_launch, 3), "pid": int(pid), "call": call}
            if "AF_UNIX" in rest:
                mm = re.search(r'sun_path=(@?"[^"]*")', rest)
                c.update(family="AF_UNIX", dst=mm.group(1) if mm else "?")
            elif "AF_INET6" in rest:
                mm = re.search(r'sin6_port=htons\((\d+)\).*?inet_pton\(AF_INET6, "([^"]+)"', rest)
                c.update(family="AF_INET6", dst=f"[{mm.group(2)}]:{mm.group(1)}" if mm else "?")
            elif "AF_INET" in rest:
                mm = re.search(r'sin_port=htons\((\d+)\), sin_addr=inet_addr\("([\d.]+)"\)', rest)
                c.update(family="AF_INET", dst=f"{mm.group(2)}:{mm.group(1)}" if mm else "?")
            else:
                continue                                    # sends on an already-connected socket
            mm = re.search(r"= (-?\d+)(?: (\w+))?", rest)
            c["result"] = (mm.group(2) or mm.group(1)) if mm else "unfinished"
            calls.append(c)
    return calls


def pkg_versions(python):
    code = ("import importlib.metadata as m, json;"
            "print(json.dumps(sorted({d.metadata['Name']: d.version for d in m.distributions()}.items())))")
    try:
        return json.loads(subprocess.run([python, "-c", code], capture_output=True, text=True, timeout=60).stdout)
    except Exception as e:
        return repr(e)


def read(path):
    try:
        with open(path) as f:
            return f.read().strip()
    except OSError:
        return None


# ---------------------------------------------------------------- main
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True)
    ap.add_argument("--rep", required=True)
    ap.add_argument("--venv", required=True)
    ap.add_argument("--mode", required=True, choices=["default", "api", "posctl"])
    ap.add_argument("--hold", type=float, default=120)
    ap.add_argument("--settle", type=float, default=3)
    ap.add_argument("--subject-env", nargs="*", default=[], help="KEY=VALUE pairs for the subject")
    ap.add_argument("--strace-dir", default="/probe/work/strace")
    ap.add_argument("--prereg-sha256", default="none")
    a = ap.parse_args()

    res = {"arm": a.arm, "rep": a.rep, "mode": a.mode, "venv": a.venv, "hold_s": a.hold,
           "prereg_sha256": a.prereg_sha256, "kernel": os.uname().release,
           "container_os": next((l.split("=", 1)[1].strip().strip('"') for l in open("/etc/os-release")
                                 if l.startswith("PRETTY_NAME=")), None),
           "harness_python": sys.version.split()[0]}
    start_listeners()
    res["selftest"] = self_test()
    if not res["selftest"]["pass"]:
        res["void"] = "self-test failed"
        json.dump(res, open(os.path.join(OUT, "result.json"), "w"), indent=1)
        return 3

    os.makedirs(HOME, exist_ok=True)
    env = {"PATH": "/usr/local/bin:/usr/bin:/bin", "HOME": HOME, "LANG": "C.UTF-8", "PYTHONDONTWRITEBYTECODE": "1"}
    env.update(kv.split("=", 1) for kv in a.subject_env)
    res["subject_env"] = env
    python = os.path.join(a.venv, "bin", "python")
    cmd = [python, "/probe/kit/subject.py", "--mode", a.mode, "--hold", str(a.hold),
           "--marks", os.path.join(OUT, "marks.json")]
    if os.path.isdir(a.strace_dir):                         # the host's strace, run by its own loader
        cmd = [f"{a.strace_dir}/ld-linux-x86-64.so.2", "--library-path", a.strace_dir, f"{a.strace_dir}/strace",
               "-f", "--seccomp-bpf", "-e", "trace=%network", "-ttt", "-s", "256",
               "-o", os.path.join(OUT, "strace.log"), "--"] + cmd
    res["subject_cmd"] = cmd

    t0_probe = time.time()
    res["files_before"] = list_files("/tmp", t0_probe)
    with open(os.path.join(OUT, "subject.stdout"), "w") as so, open(os.path.join(OUT, "subject.stderr"), "w") as se:
        t_launch = time.time()
        rc = subprocess.run(cmd, env=env, stdout=so, stderr=se).returncode
        t_exit = time.time()
    time.sleep(a.settle)
    t_cut = time.time()

    with LOCK:
        evs = [dict(e, t_rel_s=round(e["t"] - t_launch, 3)) for e in EVENTS]
    during = [e for e in evs if t_launch <= e["t"] <= t_cut]
    res.update({
        "t_launch": t_launch, "t_exit_rel_s": round(t_exit - t_launch, 3), "t_cut_rel_s": round(t_cut - t_launch, 3),
        "subject_rc": rc,
        "marks": json.load(open(os.path.join(OUT, "marks.json"))) if os.path.exists(os.path.join(OUT, "marks.json")) else None,
        "dns": [e for e in during if e["kind"] == "dns"],
        "tls": [e for e in during if e["kind"] == "tls"],
        "http": [e for e in during if e["kind"] == "http"],
        "listener_errors": [e for e in evs if e["kind"] == "dns_error"],
        "events_outside_window": [e for e in evs if e["t"] > t_cut],
        "strace": parse_strace(os.path.join(OUT, "strace.log"), t_launch),
        "files_after": list_files("/tmp", t_launch),
    })
    deviceid_raw = read(os.path.join(QUEUE_DIR, "deviceid"))
    res["deviceid"] = {"exists": deviceid_raw is not None,
                       "c_sha256_prefix8": hashlib.sha256(deviceid_raw.encode()).hexdigest().upper()[:8] if deviceid_raw else None}
    ses = read("/tmp/.ses")
    res["tmp_ses"] = {"exists": ses is not None, "content": ses}
    res["queue"] = read_queue(t_launch, deviceid_raw, read("/etc/machine-id"), socket.gethostname())
    # keep the throwaway identifiers as evidence (this container's, never the host's)
    if deviceid_raw is not None:
        shutil.copy2(os.path.join(QUEUE_DIR, "deviceid"), os.path.join(OUT, "deviceid"))
    if ses is not None:
        shutil.copy2("/tmp/.ses", os.path.join(OUT, "tmp.ses"))
    res["packages"] = pkg_versions(python)
    with open(os.path.join(OUT, "result.json"), "w") as f:
        json.dump(res, f, indent=1)
    print(json.dumps({"arm": a.arm, "rep": a.rep, "rc": rc, "dns": len(res["dns"]), "tls": len(res["tls"]),
                      "net_syscalls_inet": sum(1 for c in res["strace"] or [] if c["family"] != "AF_UNIX")}))
    return 0


if __name__ == "__main__":
    sys.exit(main())