#!/usr/bin/env python3
"""Un servidor DNS autoritativo mínimo que contesta según quién pregunta.

Es la idea de la selección por DNS de una CDN, reducida a una tabla: el prefijo de quien
consulta decide qué réplica se devuelve. Solo sabe contestar registros A, de un solo nombre.

  geo_dns.py <puerto>
"""
import ipaddress
import socket
import struct
import sys

NAME = "cdn.redes.test"
TTL = 30
# Quien consulta -> réplica más cercana. La última fila atrapa todo lo demás.
TABLE = [
    (ipaddress.ip_network("10.213.11.0/24"), "10.213.21.10"),
    (ipaddress.ip_network("10.213.12.0/24"), "10.213.22.10"),
    (ipaddress.ip_network("0.0.0.0/0"), "10.213.21.10"),
]
TYPE_A = 1
CLASS_IN = 1
FLAG_RESPONSE_AA = 0x8400
RCODE_NXDOMAIN = 3
RCODE_NOTIMP = 4


def read_question(query):
    """Devuelve (nombre, tipo, clase, fin de la pregunta). Sin compresión: es la primera pregunta."""
    labels = []
    pos = 12
    while query[pos] != 0:
        length = query[pos]
        labels.append(query[pos + 1:pos + 1 + length].decode("ascii"))
        pos += 1 + length
    pos += 1
    qtype, qclass = struct.unpack("!HH", query[pos:pos + 4])
    return ".".join(labels).lower(), qtype, qclass, pos + 4


def answer_for(source):
    address = ipaddress.ip_address(source)
    for network, replica in TABLE:
        if address in network:
            return replica
    raise ValueError("la tabla no tiene ruta por defecto")


def respond(query, source):
    """Arma la respuesta a una consulta, o None si no se puede ni leer."""
    if len(query) < 17:
        return None
    query_id = query[:2]
    try:
        name, qtype, qclass, end = read_question(query)
    except (IndexError, UnicodeDecodeError, struct.error):
        return None
    question = query[12:end]
    if name != NAME:
        return query_id + struct.pack("!HHHHH", FLAG_RESPONSE_AA | RCODE_NXDOMAIN, 1, 0, 0, 0) + question
    if qtype != TYPE_A or qclass != CLASS_IN:
        return query_id + struct.pack("!HHHHH", FLAG_RESPONSE_AA | RCODE_NOTIMP, 1, 0, 0, 0) + question
    replica = answer_for(source)
    # 0xc00c es un puntero al nombre que ya está en la pregunta.
    record = struct.pack("!HHHIH", 0xc00c, TYPE_A, CLASS_IN, TTL, 4) + socket.inet_aton(replica)
    return query_id + struct.pack("!HHHHH", FLAG_RESPONSE_AA, 1, 1, 0, 0) + question + record


def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 53
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind(("0.0.0.0", port))
    while True:
        query, peer = sock.recvfrom(512)
        reply = respond(query, peer[0])
        if reply:
            sock.sendto(reply, peer)


if __name__ == "__main__":
    main()