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]]] = []
    EXPOSED_TCP: dict[str, list[int]] = {}
    HOST_ACCESS_TCP: dict[str, list[int]] = {}
    FORWARD_TCP: list[tuple[int, str, int]] = []
    NGINX_BACKENDS: dict[str, int] = {}
    NGINX_DOMAINS: dict[str, str] = {}
    SSH_KEY: str = ""
    SSH_OPTS: str = ""
    TEST_SSH_KEY_STORE: str = ""
    SSHPIPER_TEST_USER: str = ""
    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")}

# pre-compute the set of all host-access ports (union across containers)
ALL_HOST_ACCESS_PORTS = {p for ps in HOST_ACCESS_TCP.values() for p in ps}

# published host ports
FWD_HOST_PORTS = {hp for hp, _, _ in FORWARD_TCP}


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"):
    src = enterable[0]
    dst = next(c for c in included if c != src)
    resolved = server.succeed(
        net_cmd(src, f"{DIG} +short +time=3 @{HOST_ADDR[src]} {dst}.containers")
    ).strip()
    assert resolved == ADDR[dst], f"{dst}.containers resolved to {resolved!r}, expected {ADDR[dst]}"

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

with subtest("(functional): nginx proxies to backend containers"):
    if "nginx" in enterable:
        for domain, backend in NGINX_DOMAINS.items():
            if backend in included:
                server.succeed(
                    net_cmd(
                        "nginx",
                        f"{CURL} -sfk --max-time 5 --resolve {domain}:443:{ADDR['nginx']} "
                        f"https://{domain}/ >&2",
                    )
                )

with subtest("(functional): sshpiper authenticates/routes properly"):
    if "sshpiper" in included:
        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"):
    tested = 0
    for src in enterable:
        for dst in enterable:
            if src == dst or (src, dst) in PEER_EDGES:
                continue
            ports = EXPOSED_TCP.get(dst, [])
            if not ports:
                continue
            server.fail(ct_connect(src, ADDR[dst], ports[0], timeout=2))
            tested += 1
    assert tested > 0, "no negative edge cases to test"

with subtest("(isolation): edges are one-way"):
    tested = 0
    for src, dst, tcp_ports, _ in EDGES:
        if dst in ("@host", "@internet") or src not in enterable or dst not in enterable:
            continue
        if (dst, src) in PEER_EDGES or not tcp_ports:
            continue
        server.succeed(ct_connect(src, ADDR[dst], tcp_ports[0]))
        tested += 1
        rev_ports = EXPOSED_TCP.get(src, [])
        if rev_ports:
            server.fail(ct_connect(dst, ADDR[src], rev_ports[0], timeout=2))
    assert tested > 0, "no one-way edge cases to test"

with subtest("(isolation): host access allowlist"):
    for name, ports in HOST_ACCESS_TCP.items():
        if name in enterable:
            for port in ports:
                server.succeed(ct_connect(name, HOST_ADDR[name], port))
    no_access = [c for c in enterable if c not in HOST_ACCESS_TCP]
    if no_access and ALL_HOST_ACCESS_PORTS:
        ct = no_access[0]
        for port in ALL_HOST_ACCESS_PORTS:
            server.fail(ct_connect(ct, HOST_ADDR[ct], port, timeout=2))
            # reaching the port through a different container's host addr is also blocked
            for other in HOST_ACCESS_TCP:
                if other in included:
                    server.fail(ct_connect(ct, HOST_ADDR[other], port, timeout=2))

with subtest("(isolation): host services outside the access list"):
    internet.succeed(connect(server_ip, HOST_SSHD_PORT))
    for name in enterable:
        if HOST_SSHD_PORT not in HOST_ACCESS_TCP.get(name, []):
            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"):
    src = enterable[0]
    other = enterable[1]

    before = spoof_counter()
    server.succeed(net_cmd(src, f"{PING} -c 1 -W 2 {HOST_ADDR[src]} >&2"))
    assert spoof_counter() == before, "an unspoofed packet hit the source guard"

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

with subtest("(isolation): renaming the interface changes nothing"):
    # find a container with a declared peer edge
    rename_src = rename_dst = rename_port = None
    for src, dst, tcp_ports, _ in EDGES:
        if dst in ("@host", "@internet") or src not in enterable or dst not in included:
            continue
        if tcp_ports:
            rename_src, rename_dst, rename_port = src, dst, tcp_ports[0]
            break

    assert rename_src is not None and rename_dst is not None and rename_port is not None, \
        "no peer edge to test interface rename"

    iface, host_addr = IFACE[rename_src], HOST_ADDR[rename_src]

    def rename(frm, to):
        server.succeed(net_cmd(rename_src, f"{IP} link set dev {frm} down"))
        server.succeed(net_cmd(rename_src, f"{IP} link set dev {frm} name {to}"))
        server.succeed(net_cmd(rename_src, 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(rename_src, f"{IP} route replace {route} dev {to}"))

    rename(iface, "evil0")
    try:
        server.succeed(ct_connect(rename_src, ADDR[rename_dst], rename_port))
    finally:
        rename("evil0", iface)

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

with subtest("(isolation): route bypass"):
    bypass_ct = next((c for c in enterable if c in no_gateway), None)
    if bypass_ct is None:
        bypass_ct = next((c for c in enterable if c not in allow_internet), None)
    assert bypass_ct is not None, "no container without internet to test route bypass"

    iface = IFACE[bypass_ct]
    server.succeed(
        net_cmd(bypass_ct, f"{IP} route replace default via {HOST_ADDR[bypass_ct]} dev {iface} onlink")
    )
    server.fail(net_cmd(bypass_ct, f"{CURL} -sf --max-time 5 http://{internet_ip}/ >&2"))
    unreachable = next(
        (c for c in included if c != bypass_ct and (bypass_ct, c) not in PEER_EDGES and EXPOSED_TCP.get(c, [])),
        None,
    )
    if unreachable:
        server.fail(ct_connect(bypass_ct, ADDR[unreachable], EXPOSED_TCP[unreachable][0], timeout=2))

with subtest("(isolation): no IPv6 link-local addresses"):
    for name in enterable:
        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"):
    if FWD_HOST_PORTS:
        ct = next((c for c in enterable if c not in HOST_ACCESS_TCP), enterable[0])
        for port in FWD_HOST_PORTS:
            server.fail(ct_connect(ct, HOST_ADDR[ct], port, timeout=2))
        server.fail(ct_connect(ct, server_ip, next(iter(FWD_HOST_PORTS)), timeout=2))

with subtest("(isolation): privilege escalation / container escape"):
    ct = enterable[0]
    leader = container_pid(ct)

    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.fail(ct_cmd(ct, "touch /nix/store/canary"))
    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 hp, _, _ in FORWARD_TCP:
        internet.succeed(ext_connect(hp))
    # a non-published port should be unreachable
    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}')}"
        unit = "container" if KIND[name] == "nspawn" else "microvm"
        restarts = server.succeed(f"systemctl show {unit}@{name}.service -p NRestarts --value").strip()
        assert restarts == "0", f"{unit}@{name}.service restarted {restarts} times"

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

with subtest("(functional): published ports follow the container's lifetime"):
    if FORWARD_TCP:
        hp, ct_name, _ = FORWARD_TCP[0]
        unit = "container" if KIND[ct_name] == "nspawn" else "microvm"
        server.succeed(f"systemctl stop {unit}@{ct_name}")
        internet.fail(ext_connect(hp))
        server.succeed(f"systemctl start {unit}@{ct_name}")
        server.wait_for_unit(f"{unit}@{ct_name}.service")
        internet.wait_until_succeeds(ext_connect(hp), timeout=30)

with subtest("(functional): declared devices are usable in-container"):
    for name, devices in DEVICES.items():
        for dev in devices:
            # opened, not stat'd: without DeviceAllow= the node is still visible
            server.succeed(ct_cmd(name, f"exec 3<> {dev}"))


with subtest("(functional): /dev/kvm actually initializes, not just opens"):
    # a wrong-gid or unmapped node still opens but fails the ioctls qemu
    # needs, so boot with no disk (never completes) and kill it - the
    # failure mode this catches is qemu refusing to start at all.
    for name, devices in DEVICES.items():
        if "/dev/kvm" not in devices:
            continue
        out = server.succeed(
            ct_cmd(
                name,
                f"{TIMEOUT} 3 qemu-system-x86_64 -enable-kvm -m 256 -nographic -no-reboot "
                ">/tmp/qemu.out 2>&1; cat /tmp/qemu.out",
            )
        )
        assert "Permission denied" not in out, f"{name}: qemu couldn't use /dev/kvm:\n{out}"
        assert "Could not access KVM" not in out, f"{name}: qemu couldn't use /dev/kvm:\n{out}"


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"