"""Opt-in real packet tests, using only disposable user/network namespaces.

Run: UVM_FIREWALL_INTEGRATION=1 python3 -B -m unittest discover -s tests/integration -v

Needs Linux, nft, ip, bridge, conntrack, sysctl and permission to create network
namespaces (root, or unprivileged user namespaces). No sudo, installation,
/run/netns mounts, host firewall changes or br_netfilter sysctls are used.
Bridge-connected veths stand in for guest TAPs. A third, routed namespace tests
guest egress without NAT. Ordinary test runs skip without executing commands.
"""

from __future__ import annotations

import contextlib
import json
import os
import re
import select
import shutil
import signal
import socket
import struct
import subprocess
import sys
import time
import unittest
from ipaddress import IPv4Address, IPv4Network
from pathlib import Path


class SocketEndpoint:
    """Synchronous socket operations; the kernel completes queued handshakes."""

    def __init__(self):
        self.sockets = {}

    def close(self):
        for sock in self.sockets.values():
            sock.close()

    def request(self, operation, **args):
        if operation == "run":
            result = subprocess.run(args["command"], check=True, text=True,
                                    capture_output=True, timeout=5)
            return result.stdout
        if operation in {"listen", "connect", "udp_bind"}:
            family = socket.AF_INET6 if ":" in args["ip"] else socket.AF_INET
            kind = socket.SOCK_DGRAM if operation == "udp_bind" else socket.SOCK_STREAM
            sock = socket.socket(family, kind)
            sock.settimeout(1)
            if kind == socket.SOCK_STREAM:
                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            if family == socket.AF_INET6:
                sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
            try:
                if operation == "connect":
                    sock.connect((args["ip"], args["port"]))
                else:
                    sock.bind((args["ip"], args["port"]))
                    if operation == "listen":
                        sock.listen(16)
            except TimeoutError:
                sock.close()
                return False
            except BaseException:
                sock.close()
                raise
            self.sockets[args["name"]] = sock
            return True
        if operation == "accept":
            sock, _ = self.sockets[args["listener"]].accept()
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            sock.settimeout(1)
            self.sockets[args["name"]] = sock
            return True
        if operation in {"send", "udp_send"}:
            sock = self.sockets[args["name"]]
            data = args["data"].encode("ascii")
            if operation == "send":
                sock.sendall(data)
            else:
                sock.sendto(data, (args["ip"], args["port"]))
            return True
        if operation in {"recv", "udp_recv"}:
            sock = self.sockets[args["name"]]
            sock.settimeout(0.4 if args.get("blocked") else 2)
            data = b""
            try:
                if operation == "udp_recv":
                    data, address = sock.recvfrom(4096)
                    return [data.decode("ascii"), list(address)]
                while len(data) < args["size"]:
                    chunk = sock.recv(args["size"] - len(data))
                    if not chunk:
                        break
                    data += chunk
            except TimeoutError:
                return data.decode("ascii") if data else None
            return data.decode("ascii")
        if operation == "ping":
            ident = os.getpid() & 0xFFFF
            frame = struct.pack("!BBHHH", 8, 0, 0, ident, 1) + b"uvm-private-ports!"
            padded = frame + b"\0" * (len(frame) % 2)
            checksum = sum(struct.unpack(f"!{len(padded) // 2}H", padded))
            checksum = (checksum & 65535) + (checksum >> 16)
            checksum = (~((checksum & 65535) + (checksum >> 16))) & 65535
            frame = frame[:2] + struct.pack("!H", checksum) + frame[4:]
            with socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_ICMP) as sock:
                sock.settimeout(1)
                sock.sendto(frame, (args["ip"], 0))
                deadline = time.monotonic() + 1
                while time.monotonic() < deadline:
                    try:
                        data, address = sock.recvfrom(65535)
                    except TimeoutError:
                        return False
                    offset = (data[0] & 15) * 4
                    if (address[0] == args["ip"] and data[offset] == 0
                            and data[offset + 4:offset + 8] == frame[4:8]):
                        return True
            return False
        if operation == "raw":
            with socket.socket(socket.AF_PACKET, socket.SOCK_RAW) as sock:
                sock.bind(("eth0", 0))
                sock.send(bytes.fromhex(args["frame"]))
            return True
        raise AssertionError(f"unknown endpoint operation {operation}")


def _require_isolated(parent_fd):
    parent, current = os.fstat(parent_fd), os.stat("/proc/self/ns/net")
    if (parent.st_dev, parent.st_ino) == (current.st_dev, current.st_ino):
        raise RuntimeError("refusing packet/firewall operations in the parent network namespace")
    os.close(parent_fd)


class NamespaceEndpoint:
    def __enter__(self):
        with open("/proc/self/ns/net", "rb") as parent:
            self.process = subprocess.Popen(
                ("unshare", "--net", sys.executable, "-B", str(Path(__file__).resolve()),
                 "--agent", str(parent.fileno())),
                pass_fds=(parent.fileno(),), stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE, text=True,
            )
        try:
            assert self._read() == "ready"
        except BaseException:
            self.__exit__(None, None, None)
            raise
        return self

    def __exit__(self, *_):
        if self.process.poll() is None:
            self.process.terminate()
        try:
            self.process.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self.process.kill()
            self.process.wait(timeout=3)
        for pipe in (self.process.stdin, self.process.stdout, self.process.stderr):
            pipe.close()

    def _read(self):
        if not select.select([self.process.stdout], [], [], 8)[0]:
            raise AssertionError("namespace endpoint timed out")
        line = self.process.stdout.readline()
        if not line:
            detail = self.process.stderr.read() if self.process.poll() is not None else ""
            raise AssertionError(f"namespace endpoint exited: {detail}")
        response = json.loads(line)
        if "error" in response:
            raise AssertionError(response["error"])
        return response["result"]

    def request(self, operation, **args):
        self.process.stdin.write(json.dumps({"operation": operation, **args}) + "\n")
        self.process.stdin.flush()
        return self._read()


def _agent(parent_fd):
    _require_isolated(parent_fd)
    endpoint = SocketEndpoint()
    try:
        print(json.dumps({"result": "ready"}), flush=True)
        for line in sys.stdin:
            request = json.loads(line)
            try:
                result = endpoint.request(**request)
                response = {"result": result}
            except Exception as error:
                response = {"error": f"{type(error).__name__}: {error}"}
            print(json.dumps(response), flush=True)
    finally:
        endpoint.close()


def _sandbox(parent_fd):
    _require_isolated(parent_fd)
    if not __debug__:
        raise RuntimeError("packet verification requires assertions; disable PYTHONOPTIMIZE/-O")
    # Direct script execution must import the workspace, not an installed uvm.
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
    from uvm.config import Settings
    from uvm.domain import PortRule, VmRecord
    from uvm.errors import UvmError
    from uvm.firewall import FirewallManager
    from uvm.system import CommandRunner

    runner = CommandRunner(emit=None)

    def run(*args):
        return runner.run(args, capture=True, timeout=5).stdout

    assert [link["ifname"] for link in json.loads(run("ip", "-j", "link", "show"))] == ["lo"]
    assert not any("table" in item for item in json.loads(run("nft", "--json", "list", "tables"))["nftables"])
    settings = Settings(base=Path("/tmp/uvm-firewall-namespace-test"), bridge="fw-br0",
                        network=IPv4Network("10.231.73.0/24"), gateway=IPv4Address("10.231.73.1"))
    manager = FirewallManager(settings, runner)
    manager.check_available()
    run("ip", "link", "set", "lo", "up")
    run("ip", "link", "add", settings.bridge, "type", "bridge", "stp_state", "0")
    run("ip", "addr", "add", f"{settings.gateway}/24", "dev", settings.bridge)
    run("ip", "link", "set", settings.bridge, "up")
    # This IPv4 sysctl belongs only to the freshly verified network namespace.
    run("sysctl", "-w", "net.ipv4.ip_forward=1")

    metadata = dict(cpu=1, ram_mib=128, gateway=str(settings.gateway),
                    socket="", config="", log="", network=str(settings.network), bridge=settings.bridge)
    vm = VmRecord(id="integration-vm", tap="fw-tap1", guest_ip="10.231.73.2",
                  mac="02:fc:00:00:00:01", ports=(PortRule(8080, "HTTP"),), **metadata)
    peer_vm = VmRecord(id="integration-peer", tap="fw-tap2", guest_ip="10.231.73.3",
                       mac="02:fc:00:00:00:02", ports=(PortRule(9090),), **metadata)

    with contextlib.ExitStack() as stack:
        guest = stack.enter_context(NamespaceEndpoint())
        peer = stack.enter_context(NamespaceEndpoint())
        outside = stack.enter_context(NamespaceEndpoint())
        host = SocketEndpoint()
        stack.callback(host.close)

        for endpoint, tap, address, mac in (
            (guest, vm.tap, vm.guest_ip, vm.mac),
            (peer, peer_vm.tap, peer_vm.guest_ip, peer_vm.mac),
            (outside, "fw-up", "172.30.73.2", "02:fc:00:00:00:03"),
        ):
            run("ip", "link", "add", tap, "type", "veth", "peer", "name", "fw-end")
            run("ip", "link", "set", "fw-end", "netns", str(endpoint.process.pid))
            endpoint.request("run", command=["ip", "link", "set", "fw-end", "name", "eth0"])
            endpoint.request("run", command=["ip", "link", "set", "lo", "up"])
            endpoint.request("run", command=["ip", "link", "set", "eth0", "address", mac])
            endpoint.request("run", command=["ip", "addr", "add", f"{address}/24", "dev", "eth0"])
            endpoint.request("run", command=["ip", "link", "set", "eth0", "up"])
            if endpoint is outside:
                run("ip", "addr", "add", "172.30.73.1/24", "dev", tap)
                endpoint.request("run", command=["ip", "route", "add", str(settings.network), "via", "172.30.73.1"])
                run("ip", "link", "set", tap, "up")
            else:
                record = vm if endpoint is guest else peer_vm
                run("ip", "link", "set", "dev", tap, "alias", f"uvm:{record.id}")
                run("ip", "link", "set", tap, "master", settings.bridge)
                run("ip", "link", "set", tap, "type", "bridge_slave", "group_fwd_mask", "0x4000")
                endpoint.request("run", command=["ip", "route", "add", "default", "via", str(settings.gateway)])

        manager.apply(vm)
        manager.apply(peer_vm)
        for tap in (vm.tap, peer_vm.tap):
            run("ip", "link", "set", tap, "up")
        assert manager.status(vm) == manager.status(peer_vm) == "enforced"

        def stream(client, server, target, port, name):
            assert client.request("connect", name=name, ip=target, port=port)
            assert server.request("accept", listener="listener", name=name)

        def transfer(sender, receiver, name, text, *, blocked=False):
            assert sender.request("send", name=name, data=text)
            received = receiver.request("recv", name=name, size=len(text), blocked=blocked)
            assert received == (None if blocked else text), (name, text, received, blocked)

        guest.request("listen", name="listener", ip=vm.guest_ip, port=8080)
        peer.request("listen", name="listener", ip=peer_vm.guest_ip, port=9090)
        host.request("listen", name="listener", ip=str(settings.gateway), port=9000)
        outside.request("listen", name="listener", ip="172.30.73.2", port=9000)
        stream(host, guest, vm.guest_ip, 8080, "host-service")
        stream(peer, guest, vm.guest_ip, 8080, "peer-service")
        stream(guest, host, str(settings.gateway), 9000, "host-egress")
        stream(guest, outside, "172.30.73.2", 9000, "routed-egress")
        stream(guest, peer, peer_vm.guest_ip, 9090, "peer-egress")
        pairs = ((host, guest, "host-service"), (peer, guest, "peer-service"),
                 (guest, host, "host-egress"), (guest, outside, "routed-egress"),
                 (guest, peer, "peer-egress"))
        for client, server, name in pairs:
            transfer(client, server, name, "client-before")
            transfer(server, client, name, "server-before")
        print("PASS: host/peer inbound services and host/peer/routed guest TCP egress", flush=True)

        def service_ids():
            listing = run("conntrack", "-L", "--family", "ipv4", "--proto", "tcp",
                          "--orig-dst", vm.guest_ip, "--dport", "8080", "-o", "id")
            return set(re.findall(r"\bid=(\d+)", listing))

        before = service_ids()
        assert len(before) == 2, before
        vm.ports = ()
        manager.apply(vm)
        assert manager.status(vm) == "enforced"
        assert service_ids() == before, "closing a port must preserve connection tracking"
        for client, server, name in pairs[:2]:
            transfer(client, server, name, "client-after-close", blocked=True)
            transfer(server, client, name, "server-after-close", blocked=True)
        for client, server, name in pairs[2:]:
            transfer(client, server, name, "egress-after-close")
            transfer(server, client, name, "reply-after-close")
        assert not host.request("connect", name="closed-host", ip=vm.guest_ip, port=8080)
        assert not peer.request("connect", name="closed-peer", ip=vm.guest_ip, port=8080)
        print("PASS: existing service streams drop BOTH directions, CT IDs and guest egress survive", flush=True)

        peer_vm.ports = ()
        manager.apply(peer_vm)
        transfer(guest, peer, "peer-egress", "peer-client-closed", blocked=True)
        transfer(peer, guest, "peer-egress", "peer-server-closed", blocked=True)
        transfer(guest, outside, "routed-egress", "other-table-survives")
        transfer(outside, guest, "routed-egress", "other-table-reply")

        guest.request("udp_bind", name="dns", ip=vm.guest_ip, port=40053)
        outside.request("udp_bind", name="dns", ip="172.30.73.2", port=53)
        guest.request("udp_send", name="dns", ip="172.30.73.2", port=53, data="dns-query")
        query = outside.request("udp_recv", name="dns")
        assert query[0] == "dns-query"
        outside.request("udp_send", name="dns", ip=query[1][0], port=query[1][1], data="dns-reply")
        assert guest.request("udp_recv", name="dns")[0] == "dns-reply"
        guest.request("udp_bind", name="unsolicited", ip=vm.guest_ip, port=40054)
        outside.request("udp_send", name="dns", ip=vm.guest_ip, port=40054, data="unsolicited")
        assert guest.request("udp_recv", name="unsolicited", blocked=True) is None
        for endpoint, target in ((host, vm.guest_ip), (peer, vm.guest_ip),
                                 (guest, str(settings.gateway)), (guest, "172.30.73.2")):
            assert endpoint.request("ping", ip=target), (endpoint, target)
        print("PASS: independent peer policy, outbound UDP/DNS and ICMP; unsolicited UDP blocked", flush=True)

        # Spoofed source MACs must not reach FDB learning, including the local
        # link-local path which can skip bridge prerouting entirely.
        spoof = bytes.fromhex("02aabbccddee")
        arp = struct.pack("!HHBBH", 1, 0x0800, 6, 4, 1) + spoof + socket.inet_aton(vm.guest_ip)
        arp += b"\0" * 6 + socket.inet_aton(str(settings.gateway))
        for destination in (b"\xff" * 6, bytes.fromhex("0180c200000e"), bytes.fromhex("0180c2000003")):
            guest.request("raw", frame=(destination + spoof + b"\x08\x06" + arp).hex())
        time.sleep(0.1)
        fdb = json.loads(run("bridge", "-j", "fdb", "show", "br", settings.bridge))
        assert not any(entry.get("mac") == "02:aa:bb:cc:dd:ee" for entry in fdb)

        run("ip", "-6", "addr", "add", "fd42:73::1/64", "dev", settings.bridge, "nodad")
        guest.request("run", command=["ip", "-6", "addr", "add", "fd42:73::2/64", "dev", "eth0", "nodad"])
        host_mac = json.loads(run("ip", "-j", "link", "show", "dev", settings.bridge))[0]["address"]
        run("ip", "-6", "neigh", "replace", "fd42:73::2", "lladdr", vm.mac, "nud", "permanent", "dev", settings.bridge)
        guest.request("run", command=["ip", "-6", "neigh", "replace", "fd42:73::1", "lladdr", host_mac,
                                      "nud", "permanent", "dev", "eth0"])
        guest.request("listen", name="v6-listener", ip="fd42:73::2", port=8080)
        host.request("listen", name="v6-listener", ip="fd42:73::1", port=9000)
        assert not host.request("connect", name="v6-inbound", ip="fd42:73::2", port=8080)
        assert not guest.request("connect", name="v6-outbound", ip="fd42:73::1", port=9000)
        print("PASS: pre-learning MAC identity and IPv6 blocking in both directions", flush=True)

        manager.apply(vm, quarantine=True)
        assert manager.status(vm) == "quarantined"
        transfer(guest, host, "host-egress", "quarantined-out", blocked=True)
        transfer(host, guest, "host-egress", "quarantined-in", blocked=True)
        try:
            manager.remove(vm)
        except UvmError:
            pass
        else:
            raise AssertionError("removal must refuse a live port")
        run("ip", "neigh", "flush", "dev", settings.bridge)
        assert not host.request("ping", ip=vm.guest_ip), "quarantine must also stop ARP"
        run("ip", "link", "del", vm.tap)
        manager.cleanup_connections(vm)
        manager.remove(vm)
        assert manager.status(vm) == "missing"
        assert manager.status(peer_vm) == "enforced"
        print("PASS: quarantine, offline CT cleanup and sibling-safe removal", flush=True)


@unittest.skipUnless(os.environ.get("UVM_FIREWALL_INTEGRATION") == "1",
                     "opt in with UVM_FIREWALL_INTEGRATION=1; needs isolated network namespaces")
class PrivatePortNamespaceTests(unittest.TestCase):
    def test_private_port_packet_paths_and_live_close(self):
        if sys.platform != "linux" or shutil.which("unshare") is None:
            self.skipTest("Linux and unshare are required")
        prefix = ("unshare", "--net") if os.geteuid() == 0 else (
            "unshare", "--user", "--map-root-user", "--net",
        )
        probe = subprocess.run((*prefix, "true"), text=True, capture_output=True, timeout=5)
        if probe.returncode:
            self.skipTest(f"isolated namespace permission unavailable: {probe.stderr.strip()}")
        missing = [name for name in ("ip", "bridge", "nft", "conntrack", "sysctl") if shutil.which(name) is None]
        if missing:
            self.skipTest(f"integration tools unavailable: {', '.join(missing)}")
        with open("/proc/self/ns/net", "rb") as parent:
            process = subprocess.Popen(
                (*prefix, sys.executable, "-B", str(Path(__file__).resolve()), "--sandbox", str(parent.fileno())),
                pass_fds=(parent.fileno(),), stdin=subprocess.DEVNULL,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, start_new_session=True,
            )
        try:
            stdout, stderr = process.communicate(timeout=120)
        except BaseException:
            os.killpg(process.pid, signal.SIGKILL)
            process.communicate(timeout=5)
            raise
        self.assertEqual(process.returncode, 0, f"namespace scenario failed:\n{stdout}\n{stderr}")
        print(stdout, end="")


if __name__ == "__main__":
    if len(sys.argv) == 3 and sys.argv[1] == "--sandbox":
        _sandbox(int(sys.argv[2]))
    elif len(sys.argv) == 3 and sys.argv[1] == "--agent":
        _agent(int(sys.argv[2]))
    else:
        unittest.main()