#!/usr/bin/env python3
"""Mira HTTP/2 por dentro: acepta una conexión, lista cada frame y responde por varios streams.

    python3 h2_frames.py <dirección> <puerto>

No es un servidor de HTTP/2: no valida el protocolo ni decodifica los headers (HPACK). Solo le
contesta lo justo a `curl --http2-prior-knowledge`, mandando un DATA de cada stream abierto por
turno, para que se vea la multiplexación en la captura.
"""
import socket
import struct
import sys
import threading

PREFACE = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n"
FRAME_NAMES = {0: "DATA", 1: "HEADERS", 3: "RST_STREAM", 4: "SETTINGS", 6: "PING", 7: "GOAWAY",
               8: "WINDOW_UPDATE"}
END_STREAM, ACK, END_HEADERS = 0x1, 0x1, 0x4
CHUNK = 400
CHUNKS = 5
TURN_SECONDS = 0.1
# En HPACK, 0x88 es el índice 8 de la tabla estática: ":status: 200". Un solo byte.
STATUS_200 = b"\x88"


def frame(kind, flags, stream, payload=b""):
    """Un frame: 3 bytes de largo, tipo, flags, y 4 bytes de stream (RFC 9113, sección 4.1)."""
    return struct.pack("!I", len(payload))[1:] + bytes([kind, flags]) + struct.pack("!I", stream) + payload


def read_exact(conn, count):
    data = b""
    while len(data) < count:
        chunk = conn.recv(count - len(data))
        if not chunk:
            return None
        data += chunk
    return data


def read_frame(conn):
    head = read_exact(conn, 9)
    if head is None:
        return None
    length = struct.unpack("!I", b"\0" + head[:3])[0]
    stream = struct.unpack("!I", head[5:9])[0] & 0x7FFFFFFF
    payload = read_exact(conn, length)
    return None if payload is None else (head[3], head[4], stream, payload)


def send_turn(conn, pending):
    """Manda un DATA de cada stream abierto: así se turnan dentro de la misma conexión."""
    for stream in list(pending):
        pending[stream] -= 1
        last = END_STREAM if pending[stream] == 0 else 0
        conn.sendall(frame(0, last, stream, bytes([65 + stream % 26]) * CHUNK))
        print(f"-> DATA stream={stream} largo={CHUNK}{' END_STREAM' if last else ''}", flush=True)
        if last:
            del pending[stream]


def serve(conn):
    if read_exact(conn, len(PREFACE)) != PREFACE:
        print("no empezó con el prefacio de HTTP/2", flush=True)
        return
    print("<- prefacio de HTTP/2", flush=True)
    conn.sendall(frame(4, 0, 0))
    conn.settimeout(TURN_SECONDS)
    pending = {}
    while True:
        try:
            received = read_frame(conn)
        except socket.timeout:
            # Curl no tiene más que decir por ahora: toca otro turno de datos.
            send_turn(conn, pending)
            continue
        except OSError:
            return
        if received is None:
            return
        kind, flags, stream, payload = received
        print(f"<- {FRAME_NAMES.get(kind, kind)} stream={stream} largo={len(payload)}", flush=True)
        if kind == 4 and not flags & ACK:
            conn.sendall(frame(4, ACK, 0))
        elif kind == 6 and not flags & ACK:
            conn.sendall(frame(6, ACK, 0, payload))
        elif kind == 1:
            conn.sendall(frame(1, END_HEADERS, stream, STATUS_200))
            print(f"-> HEADERS stream={stream} largo=1", flush=True)
            pending[stream] = CHUNKS
        elif kind == 7:
            return


def serve_and_close(conn):
    with conn:
        serve(conn)


def main():
    if len(sys.argv) != 3:
        sys.exit("uso: h2_frames.py <dirección> <puerto>")
    listener = socket.socket()
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind((sys.argv[1], int(sys.argv[2])))
    listener.listen(4)
    print(f"h2_frames en {sys.argv[1]}:{sys.argv[2]}", flush=True)
    while True:
        conn, _peer = listener.accept()
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        threading.Thread(target=serve_and_close, args=(conn,), daemon=True).start()


if __name__ == "__main__":
    main()