"""Exercise evaluated router rules in private user/network namespaces.

Run: python3 scripts/router-firewall-check.py /tmp/router-rules.nft
Requires ip, nft, nsenter and unshare on PATH. Never runs against the host network.
The extra NAT table models Docker publishing allowed and unlisted host ports.
"""

import os
from pathlib import Path
import subprocess as sp
import sys

original = os.environ.get("FRED_TEST_ORIGINAL_NETNS")
if not original:
    os.execvpe(
        "unshare",
        ["unshare", "-Urn", sys.executable, *sys.argv],
        {**os.environ, "FRED_TEST_ORIGINAL_NETNS": os.readlink("/proc/self/ns/net")},
    )
if os.readlink("/proc/self/ns/net") == original:
    sys.exit("Refusing to run outside a private network namespace")

children = []


def run(*args, **kw):
    return sp.run(args, check=True, text=True, **kw)


def net(pid, *args, **kw):
    return run("nsenter", "-t", str(pid), "-n", *args, **kw)


def peer(interface, address, peer_address):
    proc = sp.Popen(["unshare", "-n", "sleep", "300"])
    children.append(proc)
    # unshare must finish before moving the peer into its namespace.
    import time

    for _ in range(100):
        if os.readlink(f"/proc/{proc.pid}/ns/net") != os.readlink("/proc/self/ns/net"):
            break
        time.sleep(0.01)
    run("ip", "link", "add", interface, "type", "veth", "peer", "name", "peer0")
    run("ip", "link", "set", "peer0", "netns", str(proc.pid))
    run("ip", "addr", "add", address + "/24", "dev", interface)
    run("ip", "link", "set", interface, "up")
    net(proc.pid, "ip", "addr", "add", peer_address + "/24", "dev", "peer0")
    net(proc.pid, "ip", "link", "set", "peer0", "up")
    net(proc.pid, "ip", "link", "set", "lo", "up")
    net(proc.pid, "ip", "route", "add", "default", "via", address)
    return proc.pid


server = """
import socket,threading,time
ports=PORTS
for port in ports:
 for kind in (socket.SOCK_STREAM,socket.SOCK_DGRAM):
  s=socket.socket(socket.AF_INET,kind);s.bind(('0.0.0.0',port))
  if kind==socket.SOCK_STREAM:s.listen()
  def serve(s=s,kind=kind):
   while True:
    if kind==socket.SOCK_STREAM:
     c,a=s.accept()
     with c:
      data=c.recv(16)
      if data:c.sendall(data)
    else:
     data,a=s.recvfrom(16);s.sendto(data,a)
  threading.Thread(target=serve,daemon=True).start()
print('ready',flush=True)
time.sleep(300)
"""


def serve(pid, ports):
    cmd = [sys.executable, "-u", "-c", server.replace("PORTS", repr(ports))]
    if pid:
        cmd = ["nsenter", "-t", str(pid), "-n"] + cmd
    proc = sp.Popen(cmd, stdout=sp.PIPE, text=True)
    children.append(proc)
    assert proc.stdout.readline().strip() == "ready"


def check(pid, port, proto, expected):
    client = f"""
import socket
s=socket.socket(socket.AF_INET,socket.SOCK_{"STREAM" if proto == "tcp" else "DGRAM"})
s.settimeout(.4)
try:
 s.connect(('{"10.0.0.1" if pid == lan else "198.18.0.1"}',{port}))
 s.send(b'hello')
 ok=s.recv(16)==b'hello'
except OSError:ok=False
assert ok=={expected!r}, 'unexpected reachability: {proto}/{port}, expected {expected}'
"""
    net(pid, sys.executable, "-c", client)
    print(
        f"{'WAN' if pid == wan else 'LAN'} {proto}/{port}: {'allowed' if expected else 'blocked'}"
    )


try:
    run("ip", "link", "set", "lo", "up")
    wan = peer("eno1", "198.18.0.1", "198.18.0.2")
    lan = peer("eth0", "10.0.0.1", "10.0.0.2")
    container = peer("docker0", "172.18.0.1", "172.18.0.2")
    Path("/proc/sys/net/ipv4/ip_forward").write_text("1")
    run("nft", "-f", sys.argv[1])
    run(
        "nft",
        "-f",
        "-",
        input="""
      table ip docker_test {
        chain prerouting {
          type nat hook prerouting priority -101; policy accept;
          iifname "eno1" tcp dport { 25565, 25600, 2022 } dnat to 172.18.0.2
          iifname "eno1" udp dport { 25565, 25600, 2022 } dnat to 172.18.0.2
          iifname "eno1" tcp dport { 25564, 25601, 9999 } dnat to 172.18.0.2:9000
        }
        chain forward { type filter hook forward priority 0; policy accept; }
      }
    """,
    )
    serve(None, [22, 443, 2022])
    serve(container, [25565, 25600, 2022, 9000])
    for port in [22, 443, 25565, 25600]:
        check(wan, port, "tcp", True)
    for port in [25565, 25600]:
        check(wan, port, "udp", True)
    for port in [2022, 25564, 25601, 9999]:
        check(wan, port, "tcp", False)
    check(wan, 2022, "udp", False)
    check(lan, 2022, "tcp", True)
    print("Router namespace checks passed")
finally:
    for proc in reversed(children):
        proc.terminate()
    for proc in children:
        proc.wait()