#!/usr/bin/env python3
"""Descifra el primer vuelo de una conexión QUIC sin tener ninguna clave privada.

Las claves de los paquetes Initial salen del Destination Connection ID, que viaja a la vista,
y de una sal fija publicada en el RFC 9001. Cualquiera en el camino puede repetir la cuenta.

  quic_initial.py --vector               comprueba las claves con el ejemplo del RFC 9001
  quic_initial.py --catch nombre.prueba  hace que curl intente HTTP/3 contra un socket propio
  quic_initial.py --pcap archivo.pcap    lee los Initial del cliente en una captura

AES viene del comando openssl, y no se verifica la etiqueta de autenticación de GCM: aquí solo
interesa mostrar que el contenido se puede leer.
"""
import argparse
import hashlib
import hmac
import socket
import struct
import subprocess
import sys

INITIAL_SALT_V1 = bytes.fromhex("38762cf7f55934b34d179ae6a4c80cadccbb7f0a")
TAG_LEN = 16


def hkdf_extract(salt, ikm):
    return hmac.new(salt, ikm, hashlib.sha256).digest()


def hkdf_expand_label(secret, label, length):
    """HKDF-Expand-Label de TLS 1.3 con contexto vacío (RFC 8446, sec. 7.1)."""
    full = b"tls13 " + label.encode()
    info = struct.pack("!HB", length, len(full)) + full + b"\x00"
    output, block, counter = b"", b"", 1
    while len(output) < length:
        block = hmac.new(secret, block + info + bytes([counter]), hashlib.sha256).digest()
        output += block
        counter += 1
    return output[:length]


def client_initial_keys(dcid):
    """Devuelve (key, iv, hp): todo sale del DCID y de una sal pública."""
    initial = hkdf_extract(INITIAL_SALT_V1, dcid)
    secret = hkdf_expand_label(initial, "client in", 32)
    return (hkdf_expand_label(secret, "quic key", 16), hkdf_expand_label(secret, "quic iv", 12),
            hkdf_expand_label(secret, "quic hp", 16))


def openssl_enc(mode, key, iv, data):
    args = ["openssl", "enc", mode, "-K", key.hex(), "-nosalt"]
    if iv is not None:
        args += ["-iv", iv.hex()]
    if mode.endswith("ecb"):
        args.append("-nopad")
    return subprocess.run(args, input=data, capture_output=True, check=True).stdout


def read_varint(data, pos):
    """Enteros de largo variable de QUIC: los 2 bits altos dicen cuántos bytes ocupan."""
    length = 1 << (data[pos] >> 6)
    value = data[pos] & 0x3f
    for byte in data[pos + 1:pos + length]:
        value = value << 8 | byte
    return value, pos + length


def unprotect_initial(packet):
    """Quita la protección de un Initial del cliente y devuelve sus datos y sus frames."""
    first = packet[0]
    if first & 0xf0 != 0xc0:
        raise ValueError("no es un Initial con encabezado largo")
    version = int.from_bytes(packet[1:5], "big")
    if version != 1:
        raise ValueError(f"versión de QUIC desconocida: {version:#x}")
    pos = 5
    dcid = packet[pos + 1:pos + 1 + packet[pos]]
    pos += 1 + packet[pos]
    scid = packet[pos + 1:pos + 1 + packet[pos]]
    pos += 1 + packet[pos]
    token_len, pos = read_varint(packet, pos)
    pos += token_len
    length, pos = read_varint(packet, pos)
    key, iv, hp = client_initial_keys(dcid)

    # La muestra para la máscara empieza 4 bytes después del inicio del número de paquete.
    sample = packet[pos + 4:pos + 4 + 16]
    mask = openssl_enc("-aes-128-ecb", hp, None, sample)
    first_clear = first ^ (mask[0] & 0x0f)
    pn_len = (first_clear & 0x03) + 1
    pn_bytes = bytes(b ^ m for b, m in zip(packet[pos:pos + pn_len], mask[1:]))
    packet_number = int.from_bytes(pn_bytes, "big")

    # AES-GCM con nonce de 96 bits cifra en modo contador, y el primer bloque de datos usa el
    # contador 2. Por eso se puede descifrar con AES-CTR, si no se revisa la etiqueta.
    nonce = (int.from_bytes(iv, "big") ^ packet_number).to_bytes(12, "big")
    body = packet[pos + pn_len:pos + length - TAG_LEN]
    frames = openssl_enc("-aes-128-ctr", key, nonce + b"\x00\x00\x00\x02", body)
    return {"dcid": dcid, "scid": scid, "packet_number": packet_number, "frames": frames,
            "header_len": pos + pn_len, "size": len(packet)}


def crypto_frames(frames):
    """Recorre los frames: PADDING, PING, ACK y CRYPTO son los únicos permitidos aquí."""
    pieces = []
    pos = 0
    while pos < len(frames):
        kind = frames[pos]
        if kind == 0x00:
            pos += 1
        elif kind == 0x01:
            pos += 1
        elif kind == 0x06:
            offset, pos = read_varint(frames, pos + 1)
            length, pos = read_varint(frames, pos)
            pieces.append((offset, frames[pos:pos + length]))
            pos += length
        else:
            raise ValueError(f"frame inesperado en un Initial: {kind:#x}")
    return pieces


def reassemble(pieces):
    """Une los pedazos del flujo CRYPTO por offset: un ClientHello puede ocupar varios paquetes."""
    data = bytearray()
    for offset, piece in sorted(pieces):
        if offset > len(data):
            break
        data[offset:offset + len(piece)] = piece
    return bytes(data)


def client_hello_names(hello):
    """Devuelve (SNI, lista ALPN) de un ClientHello sin encabezado de registro."""
    if len(hello) < 4 or hello[0] != 1:
        raise ValueError("no es un ClientHello")
    pos = 4 + 2 + 32
    pos += 1 + hello[pos]
    pos += 2 + int.from_bytes(hello[pos:pos + 2], "big")
    pos += 1 + hello[pos]
    end = pos + 2 + int.from_bytes(hello[pos:pos + 2], "big")
    pos += 2
    sni, alpn = None, []
    while pos + 4 <= end:
        ext = int.from_bytes(hello[pos:pos + 2], "big")
        size = int.from_bytes(hello[pos + 2:pos + 4], "big")
        body = hello[pos + 4:pos + 4 + size]
        if ext == 0:
            sni = body[5:5 + int.from_bytes(body[3:5], "big")].decode()
        elif ext == 16:
            at = 2
            while at < len(body):
                alpn.append(body[at + 1:at + 1 + body[at]].decode())
                at += 1 + body[at]
        pos += 4 + size
    return sni, alpn


def describe(packets):
    pieces = []
    for number, packet in enumerate(packets):
        info = unprotect_initial(packet)
        found = crypto_frames(info["frames"])
        pieces += found
        print(f"datagrama {number}: {info['size']} bytes, DCID={info['dcid'].hex()}, "
              f"número de paquete {info['packet_number']}, "
              f"bytes de CRYPTO: {sum(len(piece) for _, piece in found)}")
    hello = reassemble(pieces)
    sni, alpn = client_hello_names(hello)
    print(f"ClientHello de {len(hello)} bytes")
    print(f"SNI:  {sni}")
    print(f"ALPN: {' '.join(alpn)}")
    return sni, alpn


def catch(name, wait=2.0):
    """Abre un socket UDP y le pide a curl HTTP/3 contra él. Nadie contesta, y no hace falta."""
    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
        sock.bind(("127.0.0.1", 0))
        port = sock.getsockname()[1]
        sock.settimeout(wait)
        curl = subprocess.Popen(
            ["curl", "-s", "-k", "--http3-only", "--max-time", "2", "--resolve",
             f"{name}:{port}:127.0.0.1", f"https://{name}:{port}/"])
        try:
            packets = [sock.recv(2048)]
            sock.settimeout(0.3)
            # Sale cuando vuelve a llegar el mismo paquete: es una retransmisión.
            while True:
                packet = sock.recv(2048)
                if packet[0] & 0xf0 != 0xc0 or packet in packets:
                    break
                info = unprotect_initial(packet)
                numbers = [unprotect_initial(old)["packet_number"] for old in packets]
                if info["packet_number"] in numbers:
                    break
                packets.append(packet)
        except socket.timeout:
            pass
        finally:
            curl.kill()
            curl.wait()
    return packets


def from_pcap(path):
    """Saca de una captura los datagramas Initial del cliente, con tshark."""
    done = subprocess.run(
        ["tshark", "-r", path, "-Y", "udp.dstport == 443", "-T", "fields", "-e", "udp.payload"],
        capture_output=True, text=True, check=True)
    packets = []
    for line in done.stdout.split():
        packet = bytes.fromhex(line)
        if packet[0] & 0xf0 != 0xc0:
            continue
        # Las claves salen del primer DCID. Cuando el servidor contesta, el cliente pasa a usar
        # el ID que eligió el servidor: esos paquetes ya no se abren con esta cuenta.
        if packets and packet[6:6 + packet[5]] != packets[0][6:6 + packets[0][5]]:
            break
        packets.append(packet)
    return packets


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--vector", action="store_true")
    parser.add_argument("--catch", metavar="NOMBRE")
    parser.add_argument("--pcap", metavar="ARCHIVO")
    args = parser.parse_args()
    if args.vector:
        key, iv, hp = client_initial_keys(bytes.fromhex("8394c8f03e515708"))
        print("DCID", "8394c8f03e515708")
        print("key ", key.hex())
        print("iv  ", iv.hex())
        print("hp  ", hp.hex())
    elif args.catch:
        describe(catch(args.catch))
    elif args.pcap:
        describe(from_pcap(args.pcap))
    else:
        parser.print_help()
        return 2
    return 0


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