"""Los experimentos de la semana: cada uno prepara el enlace, transfiere y captura.

    python3 scenarios.py             lista los escenarios
    python3 scenarios.py clean       corre uno; las capturas quedan en captures/

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

from netlab import CAPTURES, Lab, wait_until

LAB = Lab()
SERVER = "10.210.2.10"
ALIAS = "10.210.1.20"  # la segunda dirección de a: la de las conexiones a medio abrir
PORT = 5001
LINK = "delay 20ms rate 20mbit limit 200"  # 40 ms de RTT y 20 Mbit/s, de a hacia b
# Con una cola de 200, el arranque de una transferencia a veces la desborda. Una cola de 2000
# es mayor que todos los segmentos de 2 MB juntos: aquí nada se pierde si no lo pedimos.
DEEP_LINK = "delay 20ms rate 20mbit limit 2000"
CLASSIC = {"tcp_sack": 0, "tcp_recovery": 0, "tcp_early_retrans": 0}
MODERN = {"tcp_sack": 1, "tcp_recovery": 1, "tcp_early_retrans": 3}
VERBOSE = True


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


def report(text):
    """Muestra un resultado. Las pruebas no lo imprimen: lo revisan."""
    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.sh(node, "sh", "-c", command)
    if done.returncode != 0:
        raise RuntimeError(f"{node}: {command}\n{done.stdout}{done.stderr}")
    return done.stdout + done.stderr


def quiet(node, command):
    """Limpieza: corre un comando sin mostrarlo y sin importar si falla."""
    return LAB.sh(node, "sh", "-c", command)


def set_link(toward_b=LINK, toward_a="delay 20ms"):
    # Se borra y se vuelve a crear: con replace, netem conserva el rate anterior si no se le da
    # uno nuevo, y con rate no desordena.
    quiet("m", "tc qdisc del dev eth2 root; tc qdisc del dev eth1 root")
    run("m", f"tc qdisc add dev eth2 root netem {toward_b}")
    run("m", f"tc qdisc add dev eth1 root netem {toward_a}")


def set_tcp(values):
    """Fija sysctls de TCP en a y en b. Cada nodo tiene su propio network namespace."""
    settings = " ".join(f"net.ipv4.{name}={value}" for name, value in values.items())
    for node in "ab":
        run(node, f"sysctl -qw {settings}")


def nft_drop(rule):
    run("m", f"nft 'add rule ip w10 forward {rule} drop'")


def reset():
    """Deja todo como recién levantado: enlace bueno, sin filtros, sin programas colgados."""
    for node in "ab":
        quiet(node, "pkill sink; pkill slow_reader; pkill half_open; pkill tcpdump; "
                    "pkill blast; true")
        quiet(node, "sysctl -qw net.ipv4.tcp_sack=1 net.ipv4.tcp_recovery=1 "
                    "net.ipv4.tcp_early_retrans=3 net.ipv4.tcp_syncookies=1 "
                    "net.ipv4.tcp_syn_linear_timeouts=4")
    quiet("a", "ip link set eth1 mtu 1500; ip route change default via 10.210.1.1; "
               "ip route flush cache")
    quiet("m", f"tc qdisc del dev eth2 root; tc qdisc del dev eth1 root; "
               f"tc qdisc add dev eth2 root netem {LINK}; "
               "tc qdisc add dev eth1 root netem delay 20ms; "
               "nft delete table ip w10; nft add table ip w10; "
               "nft add chain ip w10 forward '{ type filter hook forward priority 0; }'")
    time.sleep(0.2)


class Capture:
    """Captura en a y en b mientras dura el bloque with. Quedan name-a.pcap y name-b.pcap."""

    def __init__(self, name, ports=f"tcp port {PORT}", nodes="ab"):
        self.name = name
        self.ports = ports
        self.nodes = nodes

    def path(self, node):
        return os.path.join(CAPTURES, f"{self.name}-{node}.pcap")

    def __enter__(self):
        for node in self.nodes:
            if os.path.exists(self.path(node)):
                os.remove(self.path(node))
            # -B agranda el buffer del kernel: sin eso, a toda velocidad, tcpdump pierde segmentos.
            # --immediate-mode evita perder lo último al detenerlo con una señal.
            command = (f"tcpdump -i eth1 -n -s 128 -B 16384 --immediate-mode -Z root "
                       f"-w /captures/{self.name}-{node}.pcap {self.ports}")
            say(node, command + " &")
            LAB.daemon(node, f"{command} >/tmp/tcpdump.log 2>&1")
        ready = wait_until(lambda: all("listening" in quiet(n, "cat /tmp/tcpdump.log").stdout
                                       for n in self.nodes), seconds=15)
        if not ready:
            raise RuntimeError("tcpdump no arrancó")
        time.sleep(0.3)
        return self

    def __exit__(self, *_):
        time.sleep(1.0)  # lo último todavía viene en camino por el enlace con retardo
        for node in self.nodes:
            quiet(node, "pkill -INT tcpdump")
        for node in self.nodes:
            wait_until(lambda n=node: quiet(n, "pidof tcpdump").stdout.strip() == "", seconds=10)
        return False


def start_server(args, program="sink"):
    """Arranca el servidor en b. Su salida queda en /tmp/server.out."""
    say("b", f"/w10/{program} {args} &")
    quiet("b", "rm -f /tmp/server.out")
    LAB.daemon("b", f"/w10/{program} {args} >/tmp/server.out 2>&1")
    if not wait_until(lambda: f":{args.split()[0]} " in quiet("b", "ss -ltn").stdout, seconds=10):
        raise RuntimeError("el servidor no arrancó: " + quiet("b", "cat /tmp/server.out").stdout)


def server_output(seconds=60):
    """Espera a que termine el servidor y devuelve lo que escribió."""
    wait_until(lambda: quiet("b", "pidof sink slow_reader").stdout.strip() == "",
               seconds=seconds)
    return quiet("b", "cat /tmp/server.out").stdout


def send(args=f"{SERVER} {PORT} 2000000"):
    return run("a", f"/w10/blast {args}")


def transfer(name, size=2000000, server_args="", blast_args=""):
    """Una transferencia completa dentro de una captura. Devuelve la salida del servidor."""
    with Capture(name):
        start_server(f"{PORT} {server_args}".strip())
        send(f"{SERVER} {PORT} {size} {blast_args}".strip())
        result = server_output()
    return result


def scenario_clean():
    """Transferencia de 2 MB sin pérdidas, con una cola que no se puede desbordar."""
    reset()
    set_link(DEEP_LINK)
    return transfer("clean")


def scenario_mss():
    """Baja el MTU de a a 1400: el MSS que a anuncia cambia, y también el tamaño de los datos."""
    reset()
    run("a", "ip link set eth1 mtu 1400")
    try:
        return transfer("mss", size=100000)
    finally:
        run("a", "ip link set eth1 mtu 1500")


def scenario_drop():
    """El router descarta un solo segmento, el número 40. TCP moderno: con SACK y RACK."""
    reset()
    set_link(DEEP_LINK)
    nft_drop(f"tcp dport {PORT} numgen inc mod 100000 == 40 counter")
    before = LAB.counters("a")
    result = transfer("drop")
    return result, before, LAB.counters("a")


def scenario_classic():
    """El mismo descarte, con SACK y RACK apagados: TCP con los ACK duplicados de los libros."""
    reset()
    set_link(DEEP_LINK)
    set_tcp(CLASSIC)
    nft_drop(f"tcp dport {PORT} numgen inc mod 100000 == 40 counter")
    try:
        return transfer("classic")
    finally:
        set_tcp(MODERN)


def scenario_rto():
    """Un flujo lento y un corte total de 3 s: nadie manda ACK, y queda solo el temporizador."""
    reset()
    before = LAB.counters("a")
    samples = {}
    with Capture("rto"):
        start_server(str(PORT))
        LAB.daemon("a", f"/w10/blast {SERVER} {PORT} 60000 --chunk 1000 --gap-ms 100 --nodelay")
        say("a", f"/w10/blast {SERVER} {PORT} 60000 --chunk 1000 --gap-ms 100 --nodelay &")
        time.sleep(1.5)
        samples["before"] = run("a", f"ss -tin dst {SERVER}")
        report(samples["before"])
        set_link(f"{LINK} loss 100%", "delay 20ms loss 100%")
        time.sleep(2.2)
        samples["during"] = run("a", f"ss -tin dst {SERVER}")
        report(samples["during"])
        time.sleep(0.8)
        set_link()
        server_output()
    after = LAB.counters("a")
    show_counters((before, after))
    return before, after, samples


def scenario_tail():
    """Se pierden el último segmento y el FIN: no llega ningún ACK que avise. ¿Quién lo arregla?"""
    reset()
    # Hacia el servidor van: SYN (0), ACK (1), cinco de datos (2 a 6) y el FIN (7).
    nft_drop(f"tcp dport {PORT} numgen inc mod 100000 {{ 6, 7 }} counter")
    before = LAB.counters("a")
    result = transfer("tail", size=7000)
    return result, before, LAB.counters("a")


def scenario_loss():
    """2 % de pérdida aleatoria en el sentido de los datos."""
    reset()
    set_link(f"{LINK} loss 2%")
    before = LAB.counters("a")
    result = transfer("loss", size=3000000)
    return result, before, LAB.counters("a")


def scenario_reorder():
    """Sin pérdida, sin tope de velocidad y con 3 ms de jitter: los segmentos se desordenan."""
    reset()
    set_link("delay 20ms 3ms limit 2000")
    before = LAB.counters("a")
    result = transfer("reorder", size=3000000)
    return result, before, LAB.counters("a")


def scenario_reorder_classic():
    """El mismo enlace desordenado, con el TCP de los ACK duplicados."""
    reset()
    set_tcp(CLASSIC)
    set_link("delay 20ms 3ms limit 2000")
    before = LAB.counters("a")
    try:
        result = transfer("reorder-classic", size=3000000)
    finally:
        set_tcp(MODERN)
    return result, before, LAB.counters("a")


def scenario_zero():
    """El receptor acepta y se queda 3 s sin leer: su buffer se llena y anuncia ventana cero."""
    reset()
    return transfer("zero", size=300000, server_args="--rcvbuf 16384 --stall-ms 3000")


def scenario_slow():
    """Su slow_reader del laboratorio: 2 KB cada 100 ms, o sea unos 160 kbit/s."""
    reset()
    with Capture("slow"):
        start_server(f"{PORT} 8192 2048 100", program="slow_reader")
        send(f"{SERVER} {PORT} 100000")
        result = server_output()
    return result


def scenario_limit():
    """La misma transferencia con buffers de recepción distintos: ¿quién limita?"""
    reset()
    results = {}
    for rcvbuf in (8192, 32768, 262144):
        name = f"limit-{rcvbuf}"
        results[rcvbuf] = transfer(name, size=3000000, server_args=f"--rcvbuf {rcvbuf}")
        report(f"    rcvbuf {rcvbuf}: {results[rcvbuf].strip()}")
    return results


def scenario_ports():
    """Qué devuelve connect: puerto cerrado, filtrado en silencio, rechazado, y sin ruta."""
    reset()
    results = {}
    results["closed"] = run("a", f"/w10/connect_probe {SERVER} 5999 2")
    nft_drop("tcp dport 5999")
    with Capture("filtered", ports="tcp port 5999", nodes="a"):
        results["filtered"] = run("a", f"/w10/connect_probe {SERVER} 5999 2")
    # Linux reciente reintenta el SYN cada 1 s unas veces antes de empezar a duplicar.
    run("a", "sysctl -qw net.ipv4.tcp_syn_linear_timeouts=0")
    with Capture("filtered-exp", ports="tcp port 5999", nodes="a"):
        results["filtered_exp"] = run("a", f"/w10/connect_probe {SERVER} 5999 2")
    run("a", "sysctl -qw net.ipv4.tcp_syn_linear_timeouts=4")
    reset()
    run("m", "nft 'add rule ip w10 forward tcp dport 5999 reject with icmp type admin-prohibited'")
    results["rejected"] = run("a", f"/w10/connect_probe {SERVER} 5999 2")
    reset()
    results["no_route"] = run("a", "/w10/connect_probe 10.211.0.1 5999 2")
    start_server(str(PORT))
    results["open"] = run("a", f"/w10/connect_probe {SERVER} {PORT} 2")
    server_output()
    for name, text in results.items():
        report(f"    {name}: {text.strip()}")
    return results


def half_open_round(cookies):
    """Llena la cola de SYN de un servidor con backlog 16 y prueba a conectarse."""
    reset()
    run("b", f"sysctl -qw net.ipv4.tcp_syncookies={cookies}")
    # El servidor no acepta nunca: importa la cola del handshake, no lo que pasa después.
    start_server(f"{PORT} --stay --backlog 16")
    nft_drop(f"ip saddr {SERVER} ip daddr {ALIAS} tcp flags & (syn|ack) == syn|ack")
    before = LAB.counters("b")
    LAB.daemon("a", f"/w10/half_open {ALIAS} {SERVER} {PORT} 100 12")
    say("a", f"/w10/half_open {ALIAS} {SERVER} {PORT} 100 12 &")
    time.sleep(2)
    pending = len(run("b", "ss -tn state syn-recv").splitlines()) - 1
    probe = run("a", f"/w10/connect_probe {SERVER} {PORT} 2")
    after = LAB.counters("b")
    quiet("a", "pkill half_open")
    quiet("b", "pkill sink")
    return {"pending": pending, "probe": probe, "before": before, "after": after}


def scenario_syn():
    """Una cola de SYN llena, sin y con SYN cookies."""
    results = {"off": half_open_round(0), "on": half_open_round(1)}
    for name, result in results.items():
        grew = {key: result["after"].get(key, 0) - result["before"].get(key, 0)
                for key in ("TcpExtTCPReqQFullDrop", "TcpExtSyncookiesSent", "TcpExtSyncookiesRecv")}
        report(f"    cookies {name}: {result['pending']} en SYN-RECV, {result['probe'].strip()}, {grew}")
    reset()
    return results


def scenario_delack():
    """Datos chicos cada 100 ms a un receptor que no lee: ¿cuánto espera antes del ACK?"""
    reset()
    return transfer("delack", size=4000, server_args="--stall-ms 6000",
                    blast_args="--chunk 100 --gap-ms 100 --nodelay")


def scenario_halfclose():
    """El cliente cierra su sentido y después todavía recibe la respuesta del servidor."""
    reset()
    return transfer("halfclose", size=3000, server_args="--reply", blast_args="--reply")


def scenario_timewait():
    """Se pierde el último ACK del cierre: ese es el trabajo de quien espera en TIME_WAIT."""
    reset()
    # Hacia el servidor van: SYN (0), ACK (1), FIN (2) y el último ACK (3).
    nft_drop(f"tcp dport {PORT} numgen inc mod 100000 == 3 counter")
    result = transfer("timewait", size=0)
    states = run("a", f"ss -tano state time-wait dst {SERVER}")
    report(states)
    return result, states


def close_first_round(port, flags):
    """Un servidor que cierra primero; devuelve lo que dice uno nuevo en el mismo puerto."""
    start_server(f"{port} --close-first {flags}".strip())
    run("a", f"/w10/connect_probe {SERVER} {port} 2 300")
    server_output()
    time.sleep(0.2)
    states = run("b", f"ss -tano state time-wait sport = :{port}")
    report(states)
    say("b", f"timeout 1 /w10/sink {port} {flags}")
    again = quiet("b", f"timeout 1 /w10/sink {port} {flags}")
    return states, again.stdout + again.stderr


def scenario_bindfail():
    """El servidor cierra primero y queda en TIME_WAIT: ¿puede volver a abrir ese puerto?"""
    reset()
    # Puertos propios: el TIME_WAIT de una ronda no debe estorbar a los otros escenarios.
    without = close_first_round(5003, "--no-reuseaddr")
    report(f"    sin SO_REUSEADDR: {without[1].strip() or 'bind funciona'}")
    with_reuse = close_first_round(5004, "")
    report(f"    con SO_REUSEADDR: {with_reuse[1].strip() or 'bind funciona'}")
    return without, with_reuse


def scenario_abort():
    """El cliente cierra con RST en vez de FIN; el servidor se entera en su read."""
    reset()
    with Capture("abort"):
        start_server(str(PORT))
        send(f"{SERVER} {PORT} 3000 --abort")
        result = server_output()
    return result


def scenario_closed():
    """Conectarse a un puerto cerrado: el RST del kernel de b, capturado."""
    reset()
    with Capture("closed", ports="tcp port 5999"):
        result = run("a", f"/w10/connect_probe {SERVER} 5999 2")
    return result


def scenario_rtomin():
    """Fija rto_min en la ruta de a: se ve qué parte del RTO de Linux es ese piso."""
    reset()
    rows = {}
    for floor in (None, "500ms"):
        if floor:
            run("a", f"ip route change default via 10.210.1.1 rto_min {floor}")
        start_server(str(PORT))
        LAB.daemon("a", f"/w10/blast {SERVER} {PORT} 20000 --chunk 1000 --gap-ms 100 --nodelay")
        say("a", f"/w10/blast {SERVER} {PORT} 20000 --chunk 1000 --gap-ms 100 --nodelay &")
        time.sleep(1.2)
        rows[floor or "default"] = run("a", f"ss -tin dst {SERVER}")
        server_output()
    run("a", "ip route change default via 10.210.1.1")
    for floor, text in rows.items():
        report(f"    {floor}: {' '.join(w for w in text.split() if w.startswith(('rto:', 'rtt:')))}")
    return rows


def scenario_ss():
    """Una conexión activa vista con ss cada segundo: rtt, rto, cwnd."""
    reset()
    start_server(str(PORT))
    LAB.daemon("a", f"/w10/blast {SERVER} {PORT} 8000000")
    say("a", f"/w10/blast {SERVER} {PORT} 8000000 &")
    rows = []
    for _ in range(3):
        time.sleep(0.8)
        rows.append(run("a", f"ss -tin dst {SERVER}"))
        report(rows[-1])
    server_output()
    return rows


def show_counters(result):
    """Muestra cuánto cambiaron los contadores del kernel entre (antes, después)."""
    names = ("TcpRetransSegs", "TcpExtTCPTimeouts", "TcpExtTCPLossProbes", "TcpExtTCPSackRecovery",
             "TcpExtTCPRenoRecovery")
    before, after = result
    report("  ".join(f"{name}={after.get(name, 0) - before.get(name, 0)}" for name in names))


def scenarios():
    module = sys.modules[__name__]
    return {name[len("scenario_"):]: getattr(module, name)
            for name in dir(module) if name.startswith("scenario_")}


def main():
    table = scenarios()
    if len(sys.argv) != 2 or sys.argv[1] not in table:
        width = max(len(name) for name in table)
        for name, function in table.items():
            print(f"{name:<{width}}  {function.__doc__}")
        return 2
    result = table[sys.argv[1]]()
    if isinstance(result, str):
        print(result, end="")
    elif isinstance(result, tuple) and len(result) == 3 and isinstance(result[0], str):
        print(result[0], end="")
        show_counters(result[1:])
    return 0


if __name__ == "__main__":
    sys.exit(main())
