#!/usr/bin/env python3
"""Un cache RPKI mínimo: entrega por RTR (RFC 8210) una lista fija de VRPs.

Un validador de verdad descarga los ROAs de los repositorios y comprueba sus firmas
antes de entregarlos. Este programa se salta esa parte: solo habla el lado RTR, para
ver qué hace el router con la lista. Las VRPs se leen del archivo en cada conexión.
"""
import argparse
import socket
import struct
import sys
import threading

SERIAL_QUERY, RESET_QUERY = 1, 2
CACHE_RESPONSE, IPV4_PREFIX, END_OF_DATA, CACHE_RESET, ERROR_REPORT = 3, 4, 7, 8, 10
SESSION_ID = 7007
SERIAL = 1
# Refresh, Retry y Expire, en segundos. Expire tiene un mínimo de 600 en el RFC.
REFRESH, RETRY, EXPIRE = 30, 5, 600


def read_vrps(path):
    """Cada línea: prefijo/largo maxlength as. Lo demás se ignora."""
    vrps = []
    with open(path, encoding="utf-8") as file:
        for line in file:
            line = line.split("#")[0].strip()
            if not line:
                continue
            prefix, max_length, asn = line.split()
            address, length = prefix.split("/")
            vrps.append((socket.inet_aton(address), int(length), int(max_length), int(asn)))
    return vrps


def ipv4_prefix_pdu(version, vrp):
    address, length, max_length, asn = vrp
    # Banderas = 1: anuncio.
    return struct.pack("!BBHIBBBB4sI", version, IPV4_PREFIX, 0, 20, 1, length, max_length, 0,
                       address, asn)


def end_of_data_pdu(version):
    if version == 0:
        return struct.pack("!BBHII", 0, END_OF_DATA, SESSION_ID, 12, SERIAL)
    return struct.pack("!BBHIIIII", version, END_OF_DATA, SESSION_ID, 24, SERIAL,
                       REFRESH, RETRY, EXPIRE)


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 send_table(conn, version, path):
    data = struct.pack("!BBHI", version, CACHE_RESPONSE, SESSION_ID, 8)
    for vrp in read_vrps(path):
        data += ipv4_prefix_pdu(version, vrp)
    data += end_of_data_pdu(version)
    conn.sendall(data)


def serve(conn, peer, path):
    with conn:
        while True:
            header = read_exact(conn, 8)
            if header is None:
                return
            version, kind, _, length = struct.unpack("!BBHI", header)
            if length < 8 or length > 4096:
                return
            if length > 8 and read_exact(conn, length - 8) is None:
                return
            print(f"{peer[0]}: PDU tipo {kind}, versión {version}", flush=True)
            if version > 1:
                # Código 4: versión no soportada. El router reintenta con una menor.
                conn.sendall(struct.pack("!BBHIII", 1, ERROR_REPORT, 4, 16, 0, 0))
                return
            if kind == RESET_QUERY:
                send_table(conn, version, path)
            elif kind == SERIAL_QUERY:
                # La lista es fija: no hay diferencias que entregar.
                conn.sendall(struct.pack("!BBHI", version, CACHE_RESPONSE, SESSION_ID, 8)
                             + end_of_data_pdu(version))


def main():
    parser = argparse.ArgumentParser(description="Cache RTR mínimo con una lista fija de VRPs")
    parser.add_argument("vrps")
    parser.add_argument("--port", type=int, default=8282)
    args = parser.parse_args()
    read_vrps(args.vrps)  # falla pronto si el archivo está mal
    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server.bind(("0.0.0.0", args.port))
    server.listen()
    while True:
        conn, peer = server.accept()
        threading.Thread(target=serve, args=(conn, peer, args.vrps), daemon=True).start()


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