import json
import re
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    server: Any = None
    internet: Any = None
    start_all: Any = None
    subtest: Any = None
    retry: Any = None

    NSENTER: str = ""
    CURL: str = ""
    IP: str = ""
    PING: str = ""
    DIG: str = ""
    SSH: str = ""
    OPENSSL: str = ""
    TIMEOUT: str = ""
    included: list[str] = []
    enterable: list[str] = []
    allow_internet: list[str] = []
    not_allow_internet: list[str] = []
    has_gateway: list[str] = []
    no_gateway: list[str] = []
    ADDR: dict[str, str] = {}
    HOST_ADDR: dict[str, str] = {}
    IFACE: dict[str, str] = {}
    KIND: dict[str, str] = {}
    DEVICES: dict[str, list[str]] = {}
    EDGES: list[tuple[str, str, list[int], list[int]]] = []
    SSH_KEY: str = ""
    SSH_OPTS: str = ""
    TEST_SSH_KEY_STORE: str = ""
    SSHPIPER_TEST_USER: str = ""
    HOST_EXPORTER_PORT: int = 0
    HOST_CADVISOR_PORT: int = 0
    HOST_SSHD_PORT: int = 0

SW = "/run/current-system/sw/bin"
SH = f"{SW}/sh"

PEER_EDGES = {(src, to) for src, to, _, _ in EDGES if to not in ("@host", "@internet")}


def failed_units(node_arg=""):
    out = server.succeed(f"systemctl {node_arg} list-units --state=failed --no-legend --plain || true")
    return [
        line.split()[0]
        for line in out.splitlines()
        if line.strip() and not line.split()[0].startswith("acme-")
    ]


def container_pid(name):
    return server.succeed(f"machinectl show {name} -p Leader --value").strip()


def ct_cmd(name, cmd):
    pid = container_pid(name)
    return f"{NSENTER} --mount=/proc/{pid}/ns/mnt --root=/proc/{pid}/root -- {SH} -c '{cmd}'"


def net_cmd(name, cmd):
    pid = container_pid(name)
    return f"{NSENTER} --net=/proc/{pid}/ns/net -- {cmd}"


def container_ip(name):
    out = server.succeed(net_cmd(name, f"{IP} -4 -o addr show scope global"))
    return out.split()[3].split("/")[0]


def has_ipv6_link_local(name):
    out = server.succeed(net_cmd(name, f"{IP} -6 -o addr show scope link"))
    return out.strip() != ""


def container_iface(name):
    out = server.succeed(net_cmd(name, f"{IP} -j link show"))
    for link in json.loads(out):
        if link.get("link_type") == "ether" and "LOOPBACK" not in link.get("flags", []):
            return link["ifname"]
    raise Exception(f"no ethernet link found for {name}")


def connect(host, port, timeout=5):
    """Open a TCP connection and close it again, using only tools every node has."""
    return f"{SW}/timeout {timeout} {SH} -c 'exec 3<>/dev/tcp/{host}/{port} && exec 3<&-'"


def ct_connect(name, host, port, timeout=5):
    return net_cmd(name, connect(host, port, timeout))


def ext_connect(port, timeout=5):
    """A published port, probed the way a published port is actually used.

    Probing it from a container measures a different path - containers are
    excluded from DNAT entirely.
    """
    return connect(server_ip, port, timeout)


def spoof_counter():
    out = server.succeed("nft -j list counter inet container-ip spoofed")
    return json.loads(out)["nftables"][1]["counter"]["packets"]


start_all()
server.succeed(f"install -m 600 {TEST_SSH_KEY_STORE} {SSH_KEY}")

with subtest("(functional): containers autostart on boot with no errors"):
    server.wait_for_unit("multi-user.target")
    for name in included:
        unit = "container" if KIND[name] == "nspawn" else "microvm"
        server.wait_for_unit(f"{unit}@{name}.service")
        server.succeed(f"systemctl is-active {unit}@{name}.service")

    assert failed_units() == [], f"host has failed units: {failed_units()}"
    for name in enterable:
        assert failed_units(f"-M {name}") == [], f"{name} has failed units: {failed_units(f'-M {name}')}"

internet.wait_for_unit("dummy-http.service")
internet_ip = internet.succeed("ip -4 -o addr show dev eth1 scope global").split()[3].split("/")[0]
server_ip = server.succeed("ip -4 -o addr show dev eth1 scope global").split()[3].split("/")[0]

with subtest("(functional): derived addressing matches what the ruleset assumes"):
    for name in enterable:
        assert container_ip(name) == ADDR[name], f"{name} has {container_ip(name)}, ruleset says {ADDR[name]}"
        assert container_iface(name) == IFACE[name], f"{name} is on {container_iface(name)}, ruleset says {IFACE[name]}"
    for name in [c for c in included if KIND[c] == "vm"]:
        server.succeed(f"{IP} link show {IFACE[name]}")
        server.wait_until_succeeds(f"{PING} -c 1 -W 2 {ADDR[name]}", timeout=180)

with subtest("(isolation): internet egress follows the access list"):
    for name in allow_internet:
        server.succeed(net_cmd(name, f"{CURL} -sf --max-time 5 http://{internet_ip}/ >&2"))
    for name in not_allow_internet:
        server.fail(net_cmd(name, f"{CURL} -sf --max-time 5 http://{internet_ip}/ >&2"))
    for name in has_gateway:
        server.succeed(net_cmd(name, f"{IP} route show default | grep -q ."))
    for name in no_gateway:
        server.fail(net_cmd(name, f"{IP} route show default | grep -q ."))

with subtest("(functional): container DNS still resolves"):
    resolved = server.succeed(
        net_cmd("git", f"{DIG} +short +time=3 @{HOST_ADDR['git']} nginx.containers")
    ).strip()
    assert resolved == ADDR["nginx"], f"nginx.containers resolved to {resolved!r}, expected {ADDR['nginx']}"

with subtest("(functional): web services respond"):
    for name, port in [("devenv", 8080), ("maven", 8080), ("git", 3000)]:
        server.succeed(net_cmd(name, f"{CURL} -sf --max-time 5 -o /dev/null http://{ADDR[name]}:{port}/ >&2"))

with subtest("(functional): mail accepts connections"):
    smtp_probe = net_cmd("mail", f"{SH} -c 'exec 3<>/dev/tcp/{ADDR['mail']}/25 && head -1 <&3'")
    server.wait_until_succeeds(f"{smtp_probe} | grep -q '^220'", timeout=60)

    imap_probe = net_cmd(
        "mail",
        f"{TIMEOUT} 3 {OPENSSL} s_client -connect {ADDR['mail']}:993 -quiet -verify_quiet 2>/dev/null",
    )

    def imap_ready(_):
        greeting = server.execute(f"echo | {imap_probe} | head -1")[1]
        return greeting.strip().startswith("* OK")

    retry(imap_ready, timeout_seconds=60)

with subtest("(functional): nginx proxies to backend containers"):
    server.succeed(
        net_cmd(
            "nginx",
            f"{CURL} -sfk --max-time 5 --resolve git.floweytf.com:443:{ADDR['nginx']} "
            "https://git.floweytf.com/ >&2",
        )
    )

with subtest("(functional): sshpiper authenticates/routes properly"):
    whoami = server.succeed(f"{SSH} {SSH_OPTS} -p 2222 {SSHPIPER_TEST_USER}@{ADDR['sshpiper']} whoami").strip()
    assert whoami == "runner", f"expected to land as runner, got {whoami}"
    hostname = server.succeed(
        f"{SSH} {SSH_OPTS} -p 2222 {SSHPIPER_TEST_USER}@{ADDR['sshpiper']} hostname"
    ).strip()
    assert hostname == "ssh-admin", f"expected to land on ssh-admin, got {hostname}"

with subtest("(isolation): every declared container edge connects"):
    for src, to, tcp_ports, _ in EDGES:
        if to in ("@host", "@internet") or src not in enterable or to not in included:
            continue
        for port in tcp_ports:
            server.succeed(ct_connect(src, ADDR[to], port))

with subtest("(isolation): undeclared container edges are refused"):
    for src, dst, port in [
        ("devenv", "git", 3000),
        ("git", "devenv", 8080),
        ("maven", "git", 3000),
        ("nginx", "mail", 25),
        ("ssh-botm", "ssh-admin", 22),
        ("overseer", "git", 3000),
    ]:
        assert (src, dst) not in PEER_EDGES, f"{src} -> {dst} is a declared edge, not a negative case"
        server.fail(ct_connect(src, ADDR[dst], port, timeout=2))

with subtest("(isolation): edges are one-way"):
    assert ("nginx", "git") in PEER_EDGES
    assert ("git", "nginx") not in PEER_EDGES
    server.succeed(ct_connect("nginx", ADDR["git"], 3000))
    server.fail(ct_connect("git", ADDR["nginx"], 443, timeout=2))

with subtest("(isolation): host access allowlist"):
    server.succeed(ct_connect("overseer", HOST_ADDR["overseer"], HOST_EXPORTER_PORT))
    for port in (HOST_EXPORTER_PORT, HOST_CADVISOR_PORT):
        server.fail(ct_connect("git", HOST_ADDR["git"], port, timeout=2))
        server.fail(ct_connect("git", HOST_ADDR["overseer"], port, timeout=2))

with subtest("(isolation): host services outside the access list"):
    internet.succeed(connect(server_ip, HOST_SSHD_PORT))
    for name in ("git", "overseer"):
        server.fail(ct_connect(name, HOST_ADDR[name], HOST_SSHD_PORT, timeout=2))
        server.fail(ct_connect(name, server_ip, HOST_SSHD_PORT, timeout=2))

with subtest("(isolation): source address spoofing is dropped"):
    before = spoof_counter()
    server.succeed(net_cmd("git", f"{PING} -c 1 -W 2 {HOST_ADDR['git']} >&2"))
    assert spoof_counter() == before, "an unspoofed packet hit the source guard"

    iface = IFACE["git"]
    server.succeed(net_cmd("git", f"{IP} addr add {ADDR['nginx']}/32 dev {iface}"))
    try:
        server.execute(
            net_cmd("git", f"{PING} -c 2 -W 2 -I {ADDR['nginx']} {HOST_ADDR['git']}")
        )
        assert spoof_counter() > before, "a forged source address was not dropped"
    finally:
        server.succeed(net_cmd("git", f"{IP} addr del {ADDR['nginx']}/32 dev {iface}"))

with subtest("(isolation): renaming the interface changes nothing"):
    iface, host_addr = IFACE["nginx"], HOST_ADDR["nginx"]

    def rename(frm, to):
        server.succeed(net_cmd("nginx", f"{IP} link set dev {frm} down"))
        server.succeed(net_cmd("nginx", f"{IP} link set dev {frm} name {to}"))
        server.succeed(net_cmd("nginx", f"{IP} link set dev {to} up"))
        for route in (host_addr, f"10.0.0.0/8 via {host_addr}", f"default via {host_addr}"):
            server.succeed(net_cmd("nginx", f"{IP} route replace {route} dev {to}"))

    rename(iface, "evil0")
    try:
        server.succeed(ct_connect("nginx", ADDR["git"], 3000))
        server.fail(ct_connect("nginx", ADDR["mail"], 25, timeout=2))
    finally:
        rename("evil0", iface)

with subtest("(isolation): the veth cannot be moved out of the namespace"):
    server.fail(net_cmd("git", f"{IP} link set dev {IFACE['git']} netns 1"))

with subtest("(isolation): route bypass"):
    iface = IFACE["hopper"]
    server.succeed(
        net_cmd("hopper", f"{IP} route replace default via {HOST_ADDR['hopper']} dev {iface} onlink")
    )
    server.fail(net_cmd("hopper", f"{CURL} -sf --max-time 5 http://{internet_ip}/ >&2"))
    server.fail(ct_connect("hopper", ADDR["git"], 3000, timeout=2))

with subtest("(isolation): no IPv6 link-local addresses"):
    for name in ["git", "overseer"]:
        assert not has_ipv6_link_local(name), f"{name} still has an IPv6 link-local address"

with subtest("(isolation): published ports grant no reach between containers"):
    for port in (80, 443, 25565, 22, 25):
        server.fail(ct_connect("overseer", HOST_ADDR["overseer"], port, timeout=2))
        server.fail(ct_connect("git", server_ip, port, timeout=2))

with subtest("(isolation): privilege escalation / container escape"):
    leader = container_pid("git")

    host_uid = server.succeed(f"awk '/^Uid:/{{print $2}}' /proc/{leader}/status").strip()
    assert host_uid != "0", "container's PID 1 is not idmapped away from host uid 0"

    host_pidns = server.succeed("readlink /proc/1/ns/pid").strip()
    container_pidns = server.succeed(f"readlink /proc/{leader}/ns/pid").strip()
    assert host_pidns != container_pidns, "container shares the host's PID namespace"

    server.succeed(
        f"{NSENTER} --mount=/proc/{leader}/ns/mnt --root=/proc/{leader}/root -- "
        f"{SH} -c 'echo x >> /var/lib/gitea/canary' >&2"
    )

    sshpiper_pid = container_pid("sshpiper")
    server.fail(
        f"{NSENTER} --mount=/proc/{sshpiper_pid}/ns/mnt --root=/proc/{sshpiper_pid}/root -- "
        f"{SH} -c 'echo x >> /etc/ssh/ssh_host_ed25519_key' >&2"
    )
    server.fail(
        f"{NSENTER} --mount=/proc/{leader}/ns/mnt --root=/proc/{leader}/root -- "
        f"{SH} -c 'touch /nix/store/canary' >&2"
    )
    server.fail(
        f"{NSENTER} --mount=/proc/{leader}/ns/mnt --root=/proc/{leader}/root -- "
        "/run/current-system/sw/bin/test -d /persist >&2"
    )

with subtest("(functional): external clients reach published ports"):
    for port in (80, 25565, 22):
        internet.succeed(ext_connect(port))
    internet.fail(ext_connect(8080))

with subtest("(functional): no failing-reason restarts"):
    assert failed_units() == [], f"host has failed units: {failed_units()}"
    for name in included:
        assert (
            failed_units(f"-M {name}") == []
        ), f"{name} has failed units: {failed_units(f'-M {name}')}"
        restarts = server.succeed(f"systemctl show container@{name}.service -p NRestarts --value").strip()
        assert restarts == "0", f"container@{name}.service restarted {restarts} times"

with subtest("(functional): port forwards survive an nftables reload"):
    server.succeed("systemctl reload nftables.service")
    elements = server.succeed("nft list map ip container-portforward tcp-fwd")
    assert re.search(r"\b80\s*:", elements), f"nftables reload emptied the map: {elements!r}"
    internet.wait_until_succeeds(ext_connect(80), timeout=30)

with subtest("(functional): published ports follow the container's lifetime"):
    server.succeed("systemctl stop container@nginx")
    internet.fail(ext_connect(80))
    server.succeed("systemctl start container@nginx")
    server.wait_for_unit("container@nginx.service")
    internet.wait_until_succeeds(ext_connect(80), timeout=30)

with subtest("(functional): declared devices are usable in-container"):
    for name, devices in DEVICES.items():
        for dev in devices:
            server.succeed(ct_cmd(name, f"exec 3<> {dev}"))


with subtest("(isolation): devices reach only the container that declared them"):
    for name in enterable:
        for dev in {d for ds in DEVICES.values() for d in ds} - set(DEVICES.get(name, [])):
            server.fail(ct_cmd(name, f"test -e {dev}"))


with subtest("(isolation): identity uid mapping keeps the user namespace"):
    for name in (n for n, ds in DEVICES.items() if ds):
        host_ns = server.succeed("readlink /proc/1/ns/user").strip()
        ct_ns = server.succeed(f"readlink /proc/{container_pid(name)}/ns/user").strip()
        assert host_ns != ct_ns, f"{name} shares the host user namespace"