#!/usr/bin/env python3
"""Throw-away namespace test for the CVE-2026-98432 nft mitigation.

Run only as:  unshare -rn python3 test_cve98432.py

Interfaces, sockets and nft rules all live in a private network namespace that
disappears when this process exits. The host network is never touched.
Each case makes ONE connection; nothing loops and no race is attempted.
"""
import json
import os
import socket
import subprocess
import time
import traceback

HERE = os.path.dirname(os.path.abspath(__file__))
RULES = os.path.join(HERE, "rules.nft")
DISTRO_CONF = "/etc/nftables.conf"
IP_OPTIONS = getattr(socket, "IP_OPTIONS", 4)
RR = bytes([7, 7, 4, 0, 0, 0, 0, 0])  # Record Route, 1 slot, padded to 8 bytes
TS = bytes([68, 12, 5, 0] + [0] * 8)  # Timestamp, 2 slots, 12 bytes
PLAIN_PORT, DUAL_PORT = 5000, 5001
TIMEOUT = 3.0

passed = failed = 0


def check(name, ok, detail=""):
    global passed, failed
    passed += bool(ok)
    failed += not ok
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}" + (f"  -> {detail}" if detail else ""), flush=True)


def say(msg):
    print(f"\n== {msg}", flush=True)


def run(*cmd):
    return subprocess.run(cmd, capture_output=True, text=True)


def nft_state():
    """(table names, number of rules in cve98432, drop-counter packets or None)"""
    data = json.loads(run("nft", "-j", "list", "ruleset").stdout)["nftables"]
    tables = [o["table"]["name"] for o in data if "table" in o]
    rules = [o["rule"] for o in data if "rule" in o and o["rule"]["table"] == "cve98432"]
    pkts = None
    if rules:
        pkts = next((e["counter"]["packets"] for e in rules[0]["expr"] if "counter" in e), None)
    return tables, len(rules), pkts


def listener(family, port):
    s = socket.socket(family, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    if family == socket.AF_INET6:
        s.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
        s.bind(("::", port))
    else:
        s.bind(("0.0.0.0", port))
    s.listen(8)
    return s


def attempt(srv, port, opts=None):
    """One IPv4 connect to 127.0.0.1:port.
    Returns (connected, seconds, IP options the server-side socket holds afterwards)."""
    c = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    c.settimeout(TIMEOUT)
    t0 = time.monotonic()
    try:
        if opts:
            c.setsockopt(socket.IPPROTO_IP, IP_OPTIONS, opts)
        try:
            c.connect(("127.0.0.1", port))
        except (socket.timeout, TimeoutError):
            return False, time.monotonic() - t0, None
        srv.settimeout(1.0)
        conn, _ = srv.accept()
        try:
            held = conn.getsockopt(socket.IPPROTO_IP, IP_OPTIONS, 40)
        except OSError:
            held = None
        finally:
            conn.close()
        return True, time.monotonic() - t0, held
    finally:
        c.close()


def main():
    say("0. Safety check: private, empty network namespace?")
    links = run("ip", "-o", "link").stdout.strip().splitlines()
    check("running as root inside the namespace", os.geteuid() == 0)
    check("only the loopback link exists", len(links) == 1 and ": lo:" in links[0], f"{len(links)} link(s)")
    start = run("nft", "list", "ruleset")
    check("nft ruleset is empty at start", start.returncode == 0 and start.stdout.strip() == "")
    if failed:
        print("\nAborting: not in a fresh namespace, refusing to touch nft.")
        return
    run("ip", "link", "set", "lo", "up")
    plain = listener(socket.AF_INET, PLAIN_PORT)
    dual = listener(socket.AF_INET6, DUAL_PORT)

    say("1. CONTROL - no firewall rule yet (today's unpatched kernel)")
    ok, secs, held = attempt(plain, PLAIN_PORT)
    check("plain TCP connect works, server socket holds no options", ok and held == b"", f"{secs:.2f}s held={held!r}")
    for name, opt in (("Record-Route", RR), ("Timestamp", TS)):
        ok, secs, held = attempt(plain, PLAIN_PORT, opt)
        check(f"{name} option SYN is accepted", ok, f"{secs:.2f}s")
        check(f"{name}: server child socket now holds the option (the hand-off the CVE is about)",
              bool(held), f"held={held.hex() if held else held!r}")

    say("2. Load the proposed rule file")
    r = run("nft", "-c", "-f", RULES)
    check("nft -c (syntax check) accepts the file", r.returncode == 0, r.stderr.strip())
    r1, r2 = run("nft", "-f", RULES), run("nft", "-f", RULES)
    check("loads twice without error (idempotent)", r1.returncode == 0 and r2.returncode == 0,
          (r1.stderr + r2.stderr).strip())
    _, nrules, pkts = nft_state()
    check("exactly one rule, counter at 0", nrules == 1 and pkts == 0, f"rules={nrules} packets={pkts}")
    print(run("nft", "list", "table", "inet", "cve98432").stdout)

    say("3. With the rule loaded")
    ok, secs, _ = attempt(plain, PLAIN_PORT)
    _, _, pkts = nft_state()
    check("plain TCP still connects", ok, f"{secs:.2f}s")
    check("plain TCP did not touch the counter", pkts == 0, f"packets={pkts}")
    for name, opt in (("Record-Route", RR), ("Timestamp", TS)):
        _, _, p0 = nft_state()
        ok, secs, _ = attempt(plain, PLAIN_PORT, opt)
        _, _, p1 = nft_state()
        check(f"{name} option SYN is dropped (connect times out)", not ok and secs >= TIMEOUT - 0.2, f"{secs:.2f}s")
        check(f"{name}: drop counter rose", p1 > p0, f"packets {p0} -> {p1}")

    say("4. Dual-stack listener (an IPv6 socket that also accepts IPv4)")
    ss4, ssall = run("ss", "-Hltn4").stdout, run("ss", "-Hltn").stdout
    print("  ss -Hltn4 :", ss4.strip().replace("\n", "\n              "))
    print("  ss -Hltn  :", ssall.strip().replace("\n", "\n              "))
    check("'ss -4' lists the plain IPv4 listener", f":{PLAIN_PORT}" in ss4)
    check("'ss -4' does NOT list the dual-stack listener", f":{DUAL_PORT}" not in ss4)
    check("plain 'ss' does list the dual-stack listener", f":{DUAL_PORT}" in ssall)
    ok, secs, _ = attempt(dual, DUAL_PORT)
    check("plain IPv4 client reaches the dual-stack listener", ok, f"{secs:.2f}s")
    ok, secs, _ = attempt(dual, DUAL_PORT, RR)
    check("option-bearing IPv4 SYN to the dual-stack listener is dropped", not ok, f"{secs:.2f}s")

    say("5. What the distro nftables config does (it starts with 'flush ruleset')")
    r = run("nft", "-f", DISTRO_CONF)
    tables, _, _ = nft_state()
    check("loading /etc/nftables.conf succeeds", r.returncode == 0, r.stderr.strip())
    check("our table is wiped by it", "cve98432" not in tables, f"tables now: {tables}")
    ok, secs, _ = attempt(plain, PLAIN_PORT, RR)
    check("option-bearing SYN is accepted again (exposed again)", ok, f"{secs:.2f}s")
    r = run("nft", "-f", RULES)
    ok, secs, _ = attempt(plain, PLAIN_PORT, RR)
    check("reloading our file restores the protection", r.returncode == 0 and not ok, f"{secs:.2f}s")

    say("6. Cleanup command for a proposed ExecStop")
    r1 = run("nft", "delete", "table", "inet", "cve98432")
    r2 = run("nft", "delete", "table", "inet", "cve98432")
    check("delete works while the table exists", r1.returncode == 0, r1.stderr.strip())
    check("delete errors once it is gone (so ExecStop needs a leading '-')", r2.returncode != 0, r2.stderr.strip())


if __name__ == "__main__":
    try:
        main()
    except Exception:
        traceback.print_exc()
        failed += 1
    print(f"\nRESULT: {passed} passed, {failed} failed")