#!/usr/bin/env python3
"""Servidor autoritativo mínimo: responde por una zona, por UDP y por TCP.

    python3 dns_auth.py <dirección> <archivo de zona>

Sirve para ver la jerarquía de DNS en una red propia. No es un servidor para usar en serio.
"""
import socket
import struct
import sys
import threading

import dns_wire as w

UDP_LIMIT = 512  # sin EDNS, una respuesta por UDP no pasa de aquí (RFC 1035)


def parse_rdata(rtype, fields):
    if rtype == "MX":
        return int(fields[0]), w.normalize(fields[1])
    if rtype in ("NS", "CNAME"):
        return w.normalize(fields[0])
    if rtype == "SOA":
        return (w.normalize(fields[0]), w.normalize(fields[1]), *(int(field) for field in fields[2:7]))
    if rtype == "TXT":
        text = " ".join(fields).strip('"')
        # Una forma corta de escribir un registro enorme, para ver el truncamiento.
        if text.startswith("repetir:"):
            _tag, count, unit = text.split(":", 2)
            return unit * int(count)
        return text
    return fields[0]


def load_zone(path):
    """Lee una zona: `nombre ttl tipo datos`, un registro por línea y nombres completos."""
    records = []
    with open(path, encoding="utf-8") as file:
        for line in file:
            line = line.split(";", 1)[0].strip()
            if not line:
                continue
            name, ttl, rtype, *fields = line.split()
            records.append(w.Record(w.normalize(name), w.TYPES[rtype], int(ttl), parse_rdata(rtype, fields)))
    soa = [record for record in records if record.rtype == w.TYPES["SOA"]]
    if len(soa) != 1:
        raise ValueError(f"{path}: una zona lleva exactamente un SOA")
    return soa[0].name, records


class Zone:
    def __init__(self, origin, records):
        self.origin = origin
        self.records = records
        self.soa = next(record for record in records if record.rtype == w.TYPES["SOA"])

    def find(self, name, rtype=None):
        return [record for record in self.records
                if record.name == name and (rtype is None or record.rtype == rtype)]

    def name_exists(self, name):
        # Un nombre intermedio, sin registros propios, también existe (RFC 1034, sección 4.3.2).
        return any(record.name == name or record.name.endswith("." + name) for record in self.records)

    def delegation_for(self, name):
        """Devuelve los NS del corte de zona más alto por encima de name, si hay uno."""
        labels = name.rstrip(".").split(".")
        for start in range(len(labels) - 1, -1, -1):
            candidate = ".".join(labels[start:]) + "."
            if candidate == self.origin or not w.in_bailiwick(candidate, self.origin):
                continue
            servers = self.find(candidate, w.TYPES["NS"])
            if servers:
                return servers
        return []

    def glue_for(self, hosts):
        glue = []
        for host in hosts:
            glue += self.find(host, w.TYPES["A"]) + self.find(host, w.TYPES["AAAA"])
        return glue

    def negative(self, reply, rcode):
        # El SOA va en authority: su MINIMUM limita cuánto se recuerda la respuesta negativa.
        ttl = min(self.soa.ttl, self.soa.data[6])
        reply.authority.append(w.Record(self.soa.name, self.soa.rtype, ttl, self.soa.data))
        reply.flags = (reply.flags & ~0xF) | rcode

    def answer(self, reply, name, rtype):
        if not w.in_bailiwick(name, self.origin):
            reply.flags |= w.REFUSED
            return
        servers = self.delegation_for(name)
        if servers:
            # Referral: "yo no sé, pregúntele a estos". No es autoritativa (AA = 0).
            reply.authority += servers
            reply.additional += self.glue_for(record.data for record in servers)
            return
        reply.flags |= w.FLAG_AA
        seen = set()
        while True:
            exact = self.find(name, rtype)
            if exact:
                reply.answer += exact
                break
            alias = self.find(name, w.TYPES["CNAME"])
            if alias and name not in seen:
                seen.add(name)
                reply.answer += alias
                name = alias[0].data
                if not w.in_bailiwick(name, self.origin):
                    return  # el destino es de otra zona: lo sigue el resolver
                continue
            if not reply.answer:
                self.negative(reply, w.NOERROR if self.name_exists(name) else w.NXDOMAIN)
            return
        if rtype == w.TYPES["MX"]:
            reply.additional += self.glue_for(record.data[1] for record in exact)
        elif rtype == w.TYPES["NS"]:
            reply.additional += self.glue_for(record.data for record in exact)


def respond(zone, data, limit):
    """Devuelve la respuesta a un mensaje recibido, o None si no se puede ni contestar."""
    try:
        query = w.decode(data)
    except w.WireError:
        return None, "mensaje mal formado"
    reply = w.Message(query.msg_id, w.FLAG_QR | (query.flags & w.FLAG_RD), question=query.question[:1])
    if len(query.question) != 1:
        reply.flags |= w.FORMERR
        return w.encode(reply), "sin una pregunta"
    name, rtype = query.question[0]
    zone.answer(reply, w.normalize(name), rtype)
    raw = w.encode(reply)
    summary = (f"{w.normalize(name)} {w.TYPE_NAMES.get(rtype, rtype)} -> {w.RCODE_NAMES[reply.rcode]}"
               f" an={len(reply.answer)} ns={len(reply.authority)} ar={len(reply.additional)}")
    if limit and len(raw) > limit:
        # No cabe: mandamos solo la pregunta con TC = 1, y el cliente repite por TCP.
        short = w.Message(query.msg_id, reply.flags | w.FLAG_TC, question=reply.question)
        return w.encode(short), summary + f" TC ({len(raw)} bytes)"
    return raw, summary


def log(transport, peer, msg_id, summary):
    print(f"{transport} {peer[0]}:{peer[1]} id={msg_id} {summary}", flush=True)


def serve_udp(zone, address):
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind((address, 53))
    while True:
        data, peer = sock.recvfrom(4096)
        raw, summary = respond(zone, data, UDP_LIMIT)
        msg_id = struct.unpack_from("!H", data)[0] if len(data) >= 2 else 0
        log("udp", peer, msg_id, summary)
        if raw is not None:
            sock.sendto(raw, peer)


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 serve_tcp_client(zone, conn, peer):
    # Por TCP cada mensaje lleva delante su largo en 2 bytes: un stream no conserva límites.
    with conn:
        conn.settimeout(5)
        try:
            while True:
                head = read_exact(conn, 2)
                if head is None:
                    return
                data = read_exact(conn, struct.unpack("!H", head)[0])
                if data is None:
                    return
                raw, summary = respond(zone, data, None)
                log("tcp", peer, struct.unpack_from("!H", data)[0] if len(data) >= 2 else 0, summary)
                if raw is None:
                    return
                conn.sendall(struct.pack("!H", len(raw)) + raw)
        except OSError:
            return


def serve_tcp(zone, address):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind((address, 53))
    sock.listen(16)
    while True:
        conn, peer = sock.accept()
        threading.Thread(target=serve_tcp_client, args=(zone, conn, peer), daemon=True).start()


def main():
    if len(sys.argv) != 3:
        sys.exit("uso: dns_auth.py <dirección> <archivo de zona>")
    zone = Zone(*load_zone(sys.argv[2]))
    print(f"autoritativo de {zone.origin} en {sys.argv[1]}: {len(zone.records)} registros", flush=True)
    threading.Thread(target=serve_tcp, args=(zone, sys.argv[1]), daemon=True).start()
    serve_udp(zone, sys.argv[1])


if __name__ == "__main__":
    main()