"""Los experimentos de la semana: firewall con estado, túnel WireGuard e IPsec.

    python3 scenarios.py              lista los escenarios
    python3 scenarios.py firewall     corre uno

Cada escenario imprime los comandos que corre: léanlos, porque son lo mismo que harían a mano
en el laboratorio. Hace falta la topología levantada (make up) y las claves (make).
"""
import os
import subprocess
import sys
import time

from netlab import CAPTURES, Lab, setup_tunnel, wait_until

LAB = Lab()
HERE = os.path.dirname(os.path.abspath(__file__))
VERBOSE = True


def say(node, command):
    if VERBOSE:
        print(f"  {node}# {command}")


def report(text):
    if VERBOSE:
        print(text.rstrip())


def run(node, command):
    """Corre un comando visible en un nodo; falla si el comando falla."""
    say(node, command)
    done = LAB.shell(node, command)
    if done.returncode != 0:
        raise RuntimeError(f"{node}: {command}\n{done.stdout}{done.stderr}")
    return done.stdout + done.stderr


def quiet(node, command):
    return LAB.shell(node, command)


def serve(node, addr, port=80):
    """Deja un servidor web mínimo en un nodo, sobre una dirección y un puerto."""
    LAB.daemon(node, f"cd /tmp && printf 'ok\\n' > index.html && "
                     f"python3 -m http.server {port} --bind {addr}")
    wait_until(lambda: f":{port} " in LAB.out(node, "ss", "-tln"), seconds=5)


def firewall_path():
    """El firewall que se va a cargar: la solución si está, si no el del laboratorio."""
    for folder in ("solution", "lab"):
        path = os.path.join(HERE, folder, "firewall.nft")
        if os.path.isfile(path):
            with open(path, encoding="utf-8") as file:
                if "TODO" in file.read():
                    continue
            return path
    print("Falta completar lab/firewall.nft (todavía tiene TODO).")
    sys.exit(1)


def load_firewall():
    path = firewall_path()
    say("fw", f"nft -f {os.path.basename(path)}")
    with open(path, encoding="utf-8") as file:
        done = subprocess.run(["docker", "exec", "-i", LAB.prefix + "fw", "nft", "-f", "-"],
                              input=file.read(), capture_output=True, text=True)
    if done.returncode != 0:
        raise RuntimeError(done.stdout + done.stderr)


def timed_probe(node, host, port, label):
    """Intenta conectar con nc y muestra lo medido: si conectó y cuánto tardó."""
    start = time.monotonic()
    rc = LAB.shell(node, f"nc -w3 -z {host} {port}").returncode
    elapsed = time.monotonic() - start
    report(f"    {label}: {'conecta' if rc == 0 else 'no conecta'}, {elapsed:.2f} s")


def probe(node, url, label):
    """Pide una URL con curl y muestra el código HTTP, o el error y cuánto tardó en fallar."""
    done = LAB.shell(node, f"curl -s -m4 -o /dev/null -w '%{{http_code}} %{{time_total}}' {url}")
    code, seconds = done.stdout.split()
    report(f"    {label}: " + (code if code != "000" else f"falla en {float(seconds):.2f} s"))
    return code


def ensure_tunnel():
    """Levanta el túnel si todavía no está, y espera el primer handshake."""
    if LAB.shell("fw", "ip link show wg0").returncode != 0:
        setup_tunnel(LAB)
    quiet("ofi", "ping -c1 -W2 10.215.3.10")
    wait_until(lambda: LAB.shell("fw", "wg show wg0 latest-handshakes").stdout.split("\t")[-1].strip()
               not in ("", "0"), seconds=10)


def capture(node, iface, name, expression):
    """Empieza una captura en un nodo; devuelve la función que la detiene y la lee."""
    path = f"/captures/{name}.pcap"
    LAB.daemon(node, f"tcpdump -i {iface} -n -U -w {path} {expression}")
    time.sleep(1)

    def finish(flags="-n"):
        time.sleep(1)
        quiet(node, "pkill tcpdump")
        time.sleep(0.5)
        return LAB.sh(node, "tcpdump", flags, "-r", path).stdout
    return finish


# ----------------------------------------------------------------------------------------------


def firewall():
    """El firewall con estado: de afuera solo se alcanza el servidor web de la DMZ."""
    serve("web", "100.64.15.10", 80)
    serve("ofi", "10.215.1.10", 8080)
    print("Antes del firewall (fw solo hace forwarding): afuera llega a todo.")
    probe("ext", "http://100.64.15.10/", "ext -> web (DMZ)")
    probe("ext", "http://10.215.1.10:8080/", "ext -> ofi (red interna)")
    print("\nCargamos el firewall:")
    load_firewall()
    report(run("fw", "nft list chain inet fw forward"))
    print("Después del firewall:")
    probe("ext", "http://100.64.15.10/", "ext -> web (DMZ)   ")
    probe("ext", "http://10.215.1.10:8080/", "ext -> ofi (interna)")
    probe("ofi", "http://100.64.15.10/", "ofi -> web (interna a DMZ)")
    print("\nTres resultados distintos, medidos con el tiempo de cada intento:")
    timed_probe("ext", "100.64.15.10", 443, "ext -> web:443 (aceptado por fw, sin servidor)")
    timed_probe("ext", "100.64.15.10", 22, "ext -> web:22  (sin regla, drop en forward) ")
    timed_probe("ext", "10.215.1.10", 8080, "ext -> ofi:8080 (red interna, reject)       ")


def wireguard():
    """El túnel WireGuard entre los dos sitios, y lo que ve quien está en el medio."""
    load_firewall()
    ensure_tunnel()
    report(run("fw", "wg show wg0"))
    report(run("fw", "ip link show wg0"))
    serve("b", "10.215.3.10", 80)
    print("La oficina alcanza la red del otro sitio, que no tiene ruta pública:")
    report(run("ofi", "ping -c2 -W2 10.215.3.10"))
    print("\nCapturamos el mismo ping en tres puntos:")
    LAB.daemon("fw", "tcpdump -i eth1 -n -c2 icmp -w /captures/tunel-interno-a.pcap")
    LAB.daemon("wan", "tcpdump -i eth1 -n -c2 udp port 51820 -w /captures/tunel-externo.pcap")
    LAB.daemon("b", "tcpdump -i eth1 -n -c2 icmp -w /captures/tunel-interno-b.pcap")
    time.sleep(1)
    quiet("ofi", "ping -c2 -W2 10.215.3.10")
    time.sleep(2)
    print("  en fw eth1 (interno): ICMP en claro")
    report(LAB.out("fw", "tcpdump", "-nr", "/captures/tunel-interno-a.pcap"))
    print("  en wan eth1 (externo): UDP 51820, cifrado")
    report(LAB.out("wan", "tcpdump", "-nr", "/captures/tunel-externo.pcap"))
    print("  en b eth1 (interno, destino): ICMP en claro otra vez")
    report(LAB.out("b", "tcpdump", "-nr", "/captures/tunel-interno-b.pcap"))


def ipsec():
    """IPsec a mano con ip xfrm: ESP cifra, AH solo autentica. Claves manuales, sin IKE."""
    a, b = "100.64.15.100", "100.64.15.97"  # ext y wan, enlace directo
    import binascii
    ek = "0x" + binascii.hexlify(os.urandom(20)).decode()
    ak = "0x" + binascii.hexlify(os.urandom(32)).decode()
    print("ESP en modo transporte (cifra la carga): clave manual, una SA por sentido.")
    for node in ("ext", "wan"):
        run(node, f"ip xfrm state flush; ip xfrm policy flush")
        run(node, f"ip xfrm state add src {a} dst {b} proto esp spi 0x201 mode transport "
                  f"aead 'rfc4106(gcm(aes))' {ek} 128")
        run(node, f"ip xfrm state add src {b} dst {a} proto esp spi 0x202 mode transport "
                  f"aead 'rfc4106(gcm(aes))' {ek} 128")
        for src, dst, d in ((a, b, "out"), (b, a, "in"), (b, a, "out"), (a, b, "in")):
            quiet(node, f"ip xfrm policy add src {src} dst {dst} dir {d} tmpl proto esp mode transport")
    quiet("ext", f"ping -c1 -W2 {b}")  # resuelve ARP
    LAB.daemon("wan", "tcpdump -i eth3 -nv -c2 esp -w /captures/ipsec-esp.pcap")
    time.sleep(1)
    quiet("ext", f"ping -c2 -W2 {b}")
    time.sleep(1)
    report(LAB.out("wan", "tcpdump", "-nvr", "/captures/ipsec-esp.pcap"))
    print("\nAH en modo transporte (autentica, NO cifra): el ICMP se lee, con cabecera AH.")
    for node in ("ext", "wan"):
        run(node, f"ip xfrm state flush; ip xfrm policy flush")
        run(node, f"ip xfrm state add src {a} dst {b} proto ah spi 0x301 mode transport "
                  f"auth 'hmac(sha256)' {ak}")
        run(node, f"ip xfrm state add src {b} dst {a} proto ah spi 0x302 mode transport "
                  f"auth 'hmac(sha256)' {ak}")
        for src, dst, d in ((a, b, "out"), (b, a, "in"), (b, a, "out"), (a, b, "in")):
            quiet(node, f"ip xfrm policy add src {src} dst {dst} dir {d} tmpl proto ah mode transport")
    LAB.daemon("wan", "tcpdump -i eth3 -nv -c2 -w /captures/ipsec-ah.pcap")
    time.sleep(1)
    quiet("ext", f"ping -c2 -W2 {b}")
    time.sleep(1)
    report(LAB.out("wan", "tcpdump", "-nvr", "/captures/ipsec-ah.pcap"))
    for node in ("ext", "wan"):
        quiet(node, "ip xfrm state flush; ip xfrm policy flush")


def overhead():
    """Cuánto pesa WireGuard: el tamaño del UDP externo para tres tamaños de ping."""
    ensure_tunnel()
    for data in (56, 68, 69):
        finish = capture("wan", "eth1", f"overhead-{data}", "udp port 51820")
        quiet("ofi", f"ping -c1 -W2 -s {data} 10.215.3.10")
        first = next(line for line in finish().splitlines() if "UDP" in line)
        report(f"    -s {data}: interno {data + 28}, afuera {first.split(': ', 1)[1]}")


def allowed():
    """Identidad por clave: un origen fuera de allowed-ips llega cifrado y se descarta."""
    ensure_tunnel()
    report(run("gw2", "wg show wg0 allowed-ips"))
    for source in ("100.64.15.66", "10.215.9.1"):
        tunnel = capture("gw2", "eth1", "allowed-udp", "udp port 51820")
        inner = capture("b", "eth1", "allowed-icmp", "icmp")
        result = LAB.out("fw", "ping", "-I", source, "-c2", "-W2", "10.215.3.10")
        seen_inner = len(inner().splitlines())
        seen_tunnel = len(tunnel().splitlines())
        loss = next(line for line in result.splitlines() if "packet loss" in line)
        report(f"  origen {source}: {loss}")
        report(f"    UDP del túnel en gw2: {seen_tunnel} paquetes; ICMP en b: {seen_inner} paquetes")


def pmtu():
    """PMTUD: el ICMP que le avisa a la oficina la MTU del túnel, y lo que guarda."""
    ensure_tunnel()
    quiet("ofi", "ip route flush cache")
    finish = capture("ofi", "eth1", "pmtu", "icmp")
    LAB.daemon("b", "iperf3 -s -1 -p 5301")
    wait_until(lambda: ":5301 " in LAB.out("b", "ss", "-tln"), seconds=5)
    quiet("ofi", "iperf3 -c 10.215.3.10 -p 5301 -t 2")
    report("\n".join(finish("-nvt").splitlines()[:4]))
    report(run("ofi", "ip route get 10.215.3.10"))


def cost():
    """El costo en throughput: el mismo camino de cuatro saltos, sin cifrar y por el túnel."""
    ensure_tunnel()
    run("fw", "nft flush ruleset")
    run("wan", "ip route add 10.215.3.0/24 via 100.64.15.70")
    plain = ("ip route replace 10.215.3.0/24 via 100.64.15.65", "ip route replace 10.215.1.0/24 via 100.64.15.69")
    tunnel = ("ip route replace 10.215.3.0/24 dev wg0", "ip route replace 10.215.1.0/24 dev wg0")
    try:
        for _ in range(3):
            for name, routes in (("sin cifrar", plain), ("por el túnel", tunnel)):
                quiet("fw", routes[0])
                quiet("gw2", routes[1])
                quiet("ofi", "ip route flush cache")
                LAB.daemon("b", "iperf3 -s -1 -p 5302")
                wait_until(lambda: ":5302 " in LAB.out("b", "ss", "-tln"), seconds=5)
                out = LAB.out("ofi", "iperf3", "-c", "10.215.3.10", "-p", "5302", "-t", "3", "-f", "m", timeout=30)
                line = next(row for row in out.splitlines() if "receiver" in row)
                report(f"    {name:13s} {' '.join(line.split()[6:8])}")
    finally:
        quiet("wan", "ip route del 10.215.3.0/24")
        quiet("fw", tunnel[0])
        quiet("gw2", tunnel[1])
        load_firewall()
        quiet("fw", "nft add rule inet fw forward iifname eth1 oifname wg0 accept; "
                    "nft add rule inet fw forward iifname wg0 oifname eth1 accept")


SCENARIOS = {
    "firewall": firewall,
    "wireguard": wireguard,
    "overhead": overhead,
    "allowed": allowed,
    "pmtu": pmtu,
    "cost": cost,
    "ipsec": ipsec,
}


def main():
    os.makedirs(CAPTURES, exist_ok=True)
    if len(sys.argv) != 2 or sys.argv[1] not in SCENARIOS:
        print("escenarios:")
        for name, func in SCENARIOS.items():
            print(f"  {name:12s} {func.__doc__.splitlines()[0]}")
        return
    SCENARIOS[sys.argv[1]]()


if __name__ == "__main__":
    main()