#!/usr/bin/env python3
"""Adversarial tests of the 9proc Linux layer: probe loop, debug provider,
signal machinery. Raw 9P2000 over a unix socket, plus one 9ns mount.
Usage: adv_linux_probe.py --ns <9ns> --server <9proc-demo>
Reuses the client of adv_9proc_hostile.py. Exit 1 on any failure.
"""
import argparse
import ctypes
import os
import signal
import socket
import struct
import subprocess
import sys
import tempfile
import threading
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from adv_9proc_hostile import (  # noqa: E402
    NOFID, NOTAG, OREAD, OWRITE, Nine, Rerror, Ropen, Rread, Rversion, Rwalk, Rwrite,
    Tattach, Tread, Tversion, Twrite, frame, healthy, ok, parse_stat, s16)
import adv_9proc_hostile as hostile  # noqa: E402

libc = ctypes.CDLL(None, use_errno=True)
SYS_tgkill = 234 if os.uname().machine == "x86_64" else 131  # aarch64: 131


def tgkill(pid, tid, sig):
    return libc.syscall(SYS_tgkill, pid, tid, sig)


class Srv:
    def __init__(self, server, extra=()):
        self.tmp = tempfile.mkdtemp(prefix="advlin.")
        self.path = os.path.join(self.tmp, "sock")
        self.proc = subprocess.Popen([server, "--unix", self.path, *extra], stderr=subprocess.PIPE)
        for _ in range(200):
            if os.path.exists(self.path):
                break
            time.sleep(0.02)
        self.pid = self.proc.pid

    def alive(self):
        return self.proc.poll() is None

    def stop(self):
        if self.proc.poll() is None:
            self.proc.send_signal(signal.SIGTERM)
            try:
                self.proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.proc.kill()
                self.proc.wait()
        err = self.proc.stderr.read().decode("utf-8", "replace")
        try:
            os.unlink(self.path)
        except OSError:
            pass
        try:
            os.rmdir(self.tmp)
        except OSError:
            pass
        return err


def client(path, timeout=10):
    c = Nine(path, timeout=timeout)
    c.session()
    return c


def rd(c, names, fid=50, offset=0, count=8192):
    """walk+open+one read; returns (rtype-or-tag, data-or-error-string)."""
    if c.walk_ok(0, fid, names) != len(names):
        c.clunk(fid)
        return "walkfail", None
    rt, _, rb = c.open(fid, OREAD)
    if rt != Ropen:
        c.clunk(fid)
        return "openfail", rb
    rt, d = c.read(fid, offset, count)
    c.clunk(fid)
    if rt == Rerror:
        n, = struct.unpack_from("<H", d)
        return "err", d[2:2 + n].decode()
    return rt, d


def wr(c, names, data, fid=51, offset=0):
    if c.walk_ok(0, fid, names) != len(names):
        c.clunk(fid)
        return "walkfail", None
    rt, _, rb = c.open(fid, OWRITE)
    if rt != Ropen:
        c.clunk(fid)
        return "openfail", rb
    rt, _, rb = c.write(fid, offset, data)
    c.clunk(fid)
    if rt == Rerror:
        n, = struct.unpack_from("<H", rb)
        return "err", rb[2:2 + n].decode()
    return rt, rb


def ls(c, names, fid=52):
    c.walk_ok(0, fid, names)
    c.open(fid, OREAD)
    out = c.read_all(fid) or b""
    c.clunk(fid)
    res, i = [], 0
    while i < len(out):
        n, = struct.unpack_from("<H", out, i)
        res.append(parse_stat(out[i:i + 2 + n])["name"])
        i += 2 + n
    return res


def thread_by_name(c, pid, name):
    for t in ls(c, [b"threads"]):
        if rd(c, [b"threads", t, b"name"])[1] == name and int(t) != pid:
            return t
    return None


def ticks(c):
    return int(rd(c, [b"vars", b"state", b"f", b"ticks", b"value"])[1])


def cpu_ticks(pid):
    f = open(f"/proc/{pid}/stat").read().rsplit(")", 1)[1].split()
    return int(f[11]) + int(f[12])


def fds(pid):
    return len(os.listdir(f"/proc/{pid}/fd"))


def attack_memory(server):
    print("# memory endpoints")
    s = Srv(server)
    c = client(s.path, timeout=5)
    ok("/mem/0 is an error (not a null-pointer trap on the probe thread)", rd(c, [b"mem", b"0"]) == ("err", "i/o error"))
    ok("/hex/0 open is an error", rd(c, [b"hex", b"0"])[0] == "openfail")
    ok("/mem/0 write is an error", wr(c, [b"mem", b"0"], b"x") == ("err", "i/o error"))
    ok("/mem/1 at offset 2^64-1 (wraps to 0) is an error", rd(c, [b"mem", b"1"], offset=(1 << 64) - 1) == ("err", "i/o error"))
    ok("/mem/ffffffffffff near the top of the address space is an error", rd(c, [b"mem", b"ffffffffffff"], offset=(1 << 64) - 256) == ("err", "i/o error"))
    ok("/hex/ffffffffffff is an error", rd(c, [b"hex", b"ffffffffffff"])[0] == "openfail")
    ok("49-bit address does not walk", rd(c, [b"mem", b"1000000000000"])[0] == "walkfail")
    ok("/addr/0 renders ?", rd(c, [b"addr", b"0"])[1] == b"?\n?\n?\n")
    ok("server alive after the zero/wrap probes", s.alive() and healthy(s.path))
    # /mem/maps is the whole file, however long, byte for byte
    c.walk_ok(0, 60, [b"mem", b"maps"])
    c.open(60, OREAD)
    via = c.read_all(60, 4096)
    c.clunk(60)
    real = open(f"/proc/{s.pid}/maps", "rb").read()
    ok("/mem/maps equals /proc/<pid>/maps (read in 4 KiB pieces)", via == real, f"{len(via)} vs {len(real)}")
    maps = real.decode()
    for tag in ("[stack]", "[vdso]", "[heap]"):
        m = [ln for ln in maps.splitlines() if ln.endswith(tag)]
        if not m:
            continue
        lo = int(m[0].split("-")[0], 16) + 0x100
        ok(f"/addr of {tag} renders ? (never handed to std)", rd(c, [b"addr", b"%x" % lo])[1] == b"?\n?\n?\n")
        ok(f"/hex of {tag} dumps", rd(c, [b"hex", b"%x" % lo])[0] == Rread)
    m = [ln for ln in maps.splitlines() if ln.endswith("[stack]")][0]
    hi = int(m.split()[0].split("-")[1], 16)
    rt, d = rd(c, [b"mem", b"%x" % (hi - 16)], count=4096)
    ok("read across the end of a mapping is a short read of 16 bytes", rt == Rread and len(d) == 16, (rt, d and len(d)))
    rt, d = rd(c, [b"hex", b"%x" % (hi - 16)])
    ok("hexdump across the end of a mapping stops at the boundary", rt == Rread and d.count(b"\n") == 1, (rt, d))
    code = [ln for ln in maps.splitlines() if "r-xp" in ln and "9proc-demo" in ln][0]
    clo = int(code.split("-")[0], 16)
    ok("write into read-only code is an error, not a fault", wr(c, [b"mem", b"%x" % (clo + 0x100)], b"\xcc") == ("err", "i/o error"))
    ok("server alive after memory attacks", s.alive() and healthy(s.path))
    # a big write over the demo's own globals (state onwards) must not fault the server path
    addr = rd(c, [b"vars", b"state", b"addr"])[1].decode().strip()[2:]
    # (it zeroes the demo's own globals, this connection's state included, so
    # the reply may never come; the server as a whole must keep working)
    wr(c, [b"mem", addr.encode()], b"\x00" * 65536)
    time.sleep(0.3)
    ok("server alive and serving after a 64 KiB overwrite of its own globals", s.alive() and healthy(s.path))
    s.stop()


def attack_signals(server):
    print("# signal machinery")
    s = Srv(server)
    c = client(s.path, timeout=8)
    w = thread_by_name(c, s.pid, b"worker")
    probe = thread_by_name(c, s.pid, b"9proc")
    ok("worker and probe threads found by name", bool(w and probe), (w, probe))
    names = sorted(open(f"/proc/{s.pid}/task/{t}/comm").read().strip() for t in os.listdir(f"/proc/{s.pid}/task"))
    ok("thread names are 9proc, 9proc-demo, worker", names == ["9proc", "9proc-demo", "worker"], names)

    # 4 clients hammer stacks of every thread while trap/continue interleave
    errs, count = [], [0]
    stop = threading.Event()

    def hammer(k):
        try:
            cc = client(s.path, timeout=8)
            while not stop.is_set():
                for t in (w, probe, str(s.pid).encode()):
                    rt, d = rd(cc, [b"threads", t, b"stack"], fid=10 + k)
                    count[0] += 1
                    if rt in ("walkfail", "openfail") or (rt == "err" and "i/o" not in d):
                        errs.append((t, rt, d))
        except Exception as e:  # noqa: BLE001
            errs.append(repr(e))

    ths = [threading.Thread(target=hammer, args=(k,)) for k in range(4)]
    for t in ths:
        t.start()
    rounds_ok = True
    for _ in range(5):
        wr(c, [b"runtime", b"ctl"], b"trap")
        time.sleep(0.15)
        rounds_ok &= ls(c, [b"breakpoints"]) == [w]
        rounds_ok &= b"workerLoop" in (rd(c, [b"breakpoints", w, b"stack"])[1] or b"")
        rounds_ok &= rd(c, [b"threads", w, b"stack"])[0] == Rread  # capture of a paused thread
        rounds_ok &= wr(c, [b"breakpoints", w, b"ctl"], b"continue")[0] == Rwrite
        rounds_ok &= wr(c, [b"breakpoints", w, b"ctl"], b"continue")[0] == "walkfail"  # twice: gone
    stop.set()
    for t in ths:
        t.join()
    ok("trap/inspect/continue rounds while 4 clients capture stacks", rounds_ok)
    ok(f"{count[0]} concurrent captures without a wrong answer", count[0] > 50 and not errs, errs[:3])

    # SIGTRAP from outside (tgkill, not int3): parks without corrupting the thread
    ok("tgkill SIGTRAP to the worker", tgkill(s.pid, int(w), signal.SIGTRAP) == 0)
    time.sleep(0.3)
    ok("worker listed under /breakpoints after tgkill", ls(c, [b"breakpoints"]) == [w])
    ok("its stack names workerLoop", b"workerLoop" in (rd(c, [b"breakpoints", w, b"stack"])[1] or b""))
    t1 = ticks(c)
    time.sleep(0.3)
    ok("ticks frozen while parked", ticks(c) == t1)
    ok("continue after tgkill", wr(c, [b"breakpoints", w, b"ctl"], b"continue")[0] == Rwrite)
    time.sleep(0.4)
    ok("ticks advance after continue (no instruction skipped)", ticks(c) > t1)

    # SIGTRAP on the probe thread itself: stepped over, the server keeps serving
    ok("tgkill SIGTRAP to the probe thread", tgkill(s.pid, int(probe), signal.SIGTRAP) == 0)
    time.sleep(0.2)
    ok("server serves after a SIGTRAP on its own thread", healthy(s.path))
    ok("probe thread not parked", ls(c, [b"breakpoints"]) == [])
    # process-directed SIGTRAP lands on some thread; whichever it is, it is resumable
    os.kill(s.pid, signal.SIGTRAP)
    time.sleep(0.3)
    ok("alive after kill -TRAP <pid>", s.alive() and healthy(s.path))
    for t in ls(c, [b"breakpoints"]):
        ok(f"thread {t.decode()} parked by kill -TRAP resumes", wr(c, [b"breakpoints", t, b"ctl"], b"continue")[0] == Rwrite)
    ok("continue on a never-paused tid does not walk", wr(c, [b"breakpoints", w, b"ctl"], b"continue")[0] == "walkfail")
    ok("a bogus tid does not walk", c.walk_ok(0, 31, [b"threads", b"999999"]) == 1)

    # panic: held, inspectable, capture of the held thread works, trap meanwhile harmless, continue aborts
    wr(c, [b"runtime", b"ctl"], b"panic")
    time.sleep(0.3)
    ok("panic message published", rd(c, [b"panic", b"message"])[1] == b"demo panic requested over 9p")
    ok("panic stack names workerLoop", b"workerLoop" in rd(c, [b"panic", b"stack"])[1])
    ok("capture of the held panicking thread answers", rd(c, [b"threads", w, b"stack"])[0] == Rread)
    ok("trap request while a panic is held is harmless", wr(c, [b"runtime", b"ctl"], b"trap")[0] == Rwrite and s.alive())
    ok("panic continue", wr(c, [b"panic", b"ctl"], b"continue")[0] == Rwrite)
    ok("second panic continue is an error", wr(c, [b"panic", b"ctl"], b"continue") == ("err", "file does not exist"))
    time.sleep(1.0)
    ok("process aborted after continue", not s.alive() and s.proc.poll() not in (0, None), s.proc.poll())
    s.stop()

    s = Srv(server, ["--no-hold"])
    c = client(s.path, timeout=5)
    wr(c, [b"runtime", b"ctl"], b"panic")
    time.sleep(1.0)
    ok("--no-hold: panic aborts at once", not s.alive() and s.proc.poll() not in (0, None), s.proc.poll())
    s.stop()

    s = Srv(server)
    c = client(s.path, timeout=5)
    wr(c, [b"runtime", b"ctl"], b"trap")
    time.sleep(0.3)
    ok("worker parked", ls(c, [b"breakpoints"]) != [])
    path = s.path
    s.proc.send_signal(signal.SIGTERM)
    try:
        rc = s.proc.wait(timeout=5)
    except subprocess.TimeoutExpired:
        rc = None
    ok("SIGTERM with a parked thread exits promptly", rc is not None, rc)
    ok("SIGTERM unlinks the unix socket (clean stop path)", not os.path.exists(path))
    s.stop()

    # A server killed outright leaves its socket behind. The next server
    # of the same name must recognise the corpse and take the name over:
    # the listener probes the path, and a probe that cannot tell answers
    # "live", so a caller that hands it the whole 108-byte sun_path array
    # instead of the path turns every stale socket into AlreadyListening.
    s = Srv(server)
    client(s.path, timeout=5)
    stale = s.path
    s.proc.send_signal(signal.SIGKILL)
    s.proc.wait(timeout=5)
    ok("SIGKILL leaves the socket behind", os.path.exists(stale))
    taker = subprocess.Popen([server, "--unix", stale], stderr=subprocess.PIPE)
    try:
        # The path exists throughout (the corpse, then the new socket), so
        # the connect itself is the readiness signal.
        err, c = None, None
        for _ in range(250):
            try:
                c = client(stale, timeout=10)
                break
            except OSError as e:
                err = e
                time.sleep(0.02)
        ok("a fresh server takes over a stale socket path",
           c is not None and rd(c, [b"build", b"zig_version"])[0] == Rread,
           err or taker.poll())
    except Exception as e:
        ok("a fresh server takes over a stale socket path", False, e)
    finally:
        taker.kill()
        taker.wait(timeout=5)
    if os.path.exists(stale):
        os.unlink(stale)


def attack_probe(ns, server):
    print("# probe loop and admission")
    s = Srv(server)
    c0 = cpu_ticks(s.pid)
    time.sleep(5.0)
    ok("0 CPU ticks over 5 s idle", cpu_ticks(s.pid) - c0 == 0, cpu_ticks(s.pid) - c0)
    f0 = fds(s.pid)
    for i in range(1000):
        so = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        so.connect(s.path)
        if i % 3 == 0:
            so.sendall(frame(Tversion, NOTAG, struct.pack("<I", 8192) + s16(b"9P2000")))
            so.recv(100)
        elif i % 3 == 1:
            so.sendall(b"\x10\x00\x00")  # half a frame
        so.close()
    time.sleep(1.0)
    ok("no fd leak over 1000 connect/disconnect cycles", fds(s.pid) == f0, (f0, fds(s.pid)))
    held = [client(s.path, timeout=8) for _ in range(14)]
    pl = subprocess.Popen([ns, "--unix", s.path, "--mount", "/mnt/9p", "--", "sh", "-c", "cat /mnt/9p/build/zig_version; echo; sleep 1000"],
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    seen = pl.stdout.readline().strip()
    ok("9ns mount alongside 14 attached clients", bool(seen), seen)
    time.sleep(1.0)  # past evict_idle_ms: everyone is idle now
    slow = []
    for _ in range(40):
        so = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        so.settimeout(3)
        so.connect(s.path)
        slow.append(so)
    time.sleep(1.5)
    tv = frame(Tversion, NOTAG, struct.pack("<I", 8192) + s16(b"9P2000"))
    for k in range(6):  # slowloris: one byte at a time
        for so in slow:
            try:
                so.send(tv[k:k + 1])
            except OSError:
                pass
        time.sleep(0.2)
    ok("attached clients are never evicted by a flood", all(h.path_read([b"build", b"zig_version"]) for h in held))
    ok("the 9ns mount survives the flood", pl.poll() is None)
    pl2 = subprocess.run([ns, "--unix", s.path, "--mount", "/mnt/9p", "--", "cat", "/mnt/9p/build/zig_version"], capture_output=True, timeout=30)
    ok("a new 9ns mount is refused cleanly while the table is full", pl2.returncode != 0 or bool(pl2.stdout.strip()), pl2.stderr[-100:])
    for so in slow:
        so.close()
    time.sleep(0.5)
    # output backpressure: a client that never reads must not stall or spin the loop
    nr = client(s.path, timeout=5)
    nr.walk_ok(0, 5, [b"mem", b"maps"])
    nr.open(5, OREAD)
    nr.s.settimeout(1.0)
    try:
        for i in range(300):
            nr.s.sendall(frame(Tread, i & 0xFFFE, struct.pack("<IQI", 5, 0, 65536)))
    except OSError:
        pass  # the server stopped reading it (backpressure): that is the point
    t0 = time.time()
    d = held[0].path_read([b"build", b"zig_version"])
    ok("other clients served while one never reads its replies", bool(d) and time.time() - t0 < 1.0, f"{time.time() - t0:.2f}s")
    c0 = cpu_ticks(s.pid)
    time.sleep(2.0)
    ok("no spin with pending output on a stalled client (<= 2 ticks in 2 s)", cpu_ticks(s.pid) - c0 <= 2, cpu_ticks(s.pid) - c0)
    nr.close()
    time.sleep(0.3)
    # two clients sleep at once: a third is still served, both sleepers get their answer on time
    for h in held[2:]:
        h.close()
    time.sleep(0.3)

    def sleeper(ms):
        a = client(s.path, timeout=8)
        a.walk_ok(0, 1, [b"runtime", b"ctl"])
        a.open(1, OWRITE)
        cmd = b"sleep-ms %d" % ms
        a.raw(frame(Twrite, 3, struct.pack("<IQI", 1, 0, len(cmd)) + cmd))
        return a

    t0 = time.time()
    a = sleeper(1500)
    b = sleeper(1500)
    d = held[1].path_read([b"build", b"zig_version"])
    ok("third client served during two concurrent sleep-ms", bool(d) and time.time() - t0 < 1.0, f"{time.time() - t0:.2f}s")
    ra = a.recv_frame()[0]
    rb = b.recv_frame()[0]
    dt = time.time() - t0
    ok("both sleepers answered after ~1.5 s, not serialized", ra == Rwrite and rb == Rwrite and dt < 2.5, f"{dt:.2f}s")
    a.close()
    b.close()
    ok("healthy after probe attacks", healthy(s.path))
    pl.terminate()
    pl.wait(timeout=10)
    s.stop()

    # --stdio: EOF ends the process with exit 0 and no allocator leak report
    p = subprocess.Popen([server, "--stdio"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    p.stdin.write(frame(Tversion, NOTAG, struct.pack("<I", 8192) + s16(b"9P2000")))
    p.stdin.flush()
    h = p.stdout.read(4)
    n, = struct.unpack("<I", h)
    body = p.stdout.read(n - 4)
    ok("--stdio answers Tversion", body[0] == Rversion)
    p.stdin.close()
    try:
        rc = p.wait(timeout=5)
    except subprocess.TimeoutExpired:
        p.kill()
        rc = None
    err = p.stderr.read().decode("utf-8", "replace")
    ok("--stdio exits 0 on EOF", rc == 0, rc)
    ok("no leak report or stack trace on stderr at exit", "leaked" not in err and "in _start" not in err, err[-300:])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ns", required=True)
    ap.add_argument("--server", required=True)
    args = ap.parse_args()
    attack_memory(args.server)
    attack_signals(args.server)
    if subprocess.run(["unshare", "-Urm", "true"], capture_output=True).returncode == 0 and os.path.exists("/dev/fuse"):
        attack_probe(args.ns, args.server)
    else:
        print("# probe/admission: SKIP (namespaces or /dev/fuse unavailable)")
    print(f"# {hostile.PASSES} passed, {len(hostile.FAILS)} failed")
    sys.exit(1 if hostile.FAILS else 0)


if __name__ == "__main__":
    main()