#!/usr/bin/env python3
"""Resolver recursivo mínimo, con cache: parte de la raíz y baja por las delegaciones.

    python3 dns_resolver.py <dirección> <dirección de la raíz>
"""
import math
import secrets
import socket
import struct
import sys
import threading
import time

import dns_wire as w

MAX_STEPS = 24
QUERY_TIMEOUT = 1.0
UDP_LIMIT = 512


def acceptable_reply(reply, msg_id, question, peer, server):
    """Solo se acepta una respuesta que coincide en servidor, ID y pregunta.

    Quien quiera colar una respuesta falsa tiene que adivinar todo esto, incluido el puerto,
    que aquí revisa el kernel porque el socket está conectado al servidor.
    """
    if peer[0] != server or not reply.flags & w.FLAG_QR or reply.msg_id != msg_id:
        return False
    if len(reply.question) != 1:
        return False
    name, rtype = reply.question[0]
    return w.normalize(name) == w.normalize(question[0]) and rtype == question[1]


def usable_glue(reply, zone):
    """Se guardan solo las direcciones de nombres dentro de la zona delegada."""
    return [record for record in reply.additional
            if record.rtype in (w.TYPES["A"], w.TYPES["AAAA"]) and w.in_bailiwick(record.name, zone)]


class Cache:
    """Guarda conjuntos de registros hasta que vence su TTL."""

    def __init__(self, clock=time.monotonic):
        self.clock = clock
        self.entries = {}
        self.lock = threading.Lock()

    def put(self, name, rtype, records, ttl=None, rcode=w.NOERROR):
        ttl = min((record.ttl for record in records), default=0) if ttl is None else ttl
        if ttl > 0:
            with self.lock:
                self.entries[(name, rtype)] = (self.clock() + ttl, records, rcode)

    def get(self, name, rtype):
        """Devuelve (registros con el TTL que les queda, rcode), o None si no está o venció."""
        with self.lock:
            entry = self.entries.get((name, rtype))
            if entry is None:
                return None
            expires, records, rcode = entry
            now = self.clock()
            if expires <= now:
                del self.entries[(name, rtype)]
                return None
            left = math.ceil(expires - now)
        return [w.Record(r.name, r.rtype, left, r.data) for r in records], rcode


class Resolver:
    def __init__(self, root):
        self.root = root
        self.cache = Cache()

    def ask(self, server, name, rtype):
        """Una pregunta a un servidor, con ID nuevo y un socket (y puerto) nuevo cada vez."""
        msg_id = secrets.randbelow(65536)
        data = w.encode(w.Message(msg_id, 0, question=[(name, rtype)]))
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            sock.settimeout(QUERY_TIMEOUT)
            sock.connect((server, 53))
            sock.send(data)
            raw = sock.recv(4096)
        reply = w.decode(raw)
        if not acceptable_reply(reply, msg_id, (name, rtype), (server, 53), server):
            raise w.WireError("respuesta que no corresponde a la pregunta")
        if reply.flags & w.FLAG_TC:
            reply = self.ask_tcp(server, name, rtype)
        return reply

    def ask_tcp(self, server, name, rtype):
        msg_id = secrets.randbelow(65536)
        data = w.encode(w.Message(msg_id, 0, question=[(name, rtype)]))
        with socket.create_connection((server, 53), timeout=QUERY_TIMEOUT * 2) as conn:
            conn.sendall(struct.pack("!H", len(data)) + data)
            with conn.makefile("rb") as stream:
                head = stream.read(2)
                if len(head) != 2:
                    raise w.WireError("TCP cerrado antes de la respuesta")
                raw = stream.read(struct.unpack("!H", head)[0])
        reply = w.decode(raw)
        if not acceptable_reply(reply, msg_id, (name, rtype), (server, 53), server):
            raise w.WireError("respuesta TCP que no corresponde a la pregunta")
        return reply

    def closest_servers(self, name):
        """Busca en cache los NS del ancestro más cercano. Sin ninguno, la raíz."""
        labels = name.rstrip(".").split(".")
        for start in range(len(labels)):
            zone = ".".join(labels[start:]) + "."
            cached = self.cache.get(zone, w.TYPES["NS"])
            if not cached:
                continue
            addresses = []
            for record in cached[0]:
                glue = self.cache.get(record.data, w.TYPES["A"])
                if glue:
                    addresses += [r.data for r in glue[0]]
            if addresses:
                return addresses
        return [self.root]

    def resolve(self, name, rtype, depth=0):
        """Devuelve (rcode, registros, authority). Sigue referrals y CNAME."""
        name = w.normalize(name)
        chain = []
        for _ in range(MAX_STEPS):
            cached = self.cache.get(name, rtype)
            if cached:
                records, rcode = cached
                return rcode, chain + records, []
            alias = self.cache.get(name, w.TYPES["CNAME"])
            if alias:
                chain += alias[0]
                name = alias[0][0].data
                continue
            step = self.iterate(name, rtype, depth)
            if step is None:
                return w.SERVFAIL, chain, []
            reply = step
            if reply.rcode == w.NXDOMAIN or (reply.rcode == w.NOERROR and not reply.answer):
                self.remember_negative(name, rtype, reply)
                return reply.rcode, chain, reply.authority
            if reply.rcode != w.NOERROR:
                return w.SERVFAIL, chain, []
            self.remember_answer(reply)
        return w.SERVFAIL, chain, []

    def iterate(self, name, rtype, depth):
        """Pregunta bajando por la jerarquía hasta tener una respuesta final."""
        servers = self.closest_servers(name)
        for _ in range(MAX_STEPS):
            reply = None
            for server in servers:
                try:
                    reply = self.ask(server, name, rtype)
                    break
                except (OSError, w.WireError) as error:
                    print(f"  {server}: {error}", flush=True)
            if reply is None:
                return None
            referral = [r for r in reply.authority if r.rtype == w.TYPES["NS"]]
            if reply.answer or reply.flags & w.FLAG_AA or not referral or reply.rcode != w.NOERROR:
                return reply
            servers = self.follow_referral(reply, referral, depth)
            if not servers:
                return None
        return None

    def follow_referral(self, reply, referral, depth):
        zone = referral[0].name
        self.cache.put(zone, w.TYPES["NS"], referral)
        glue = usable_glue(reply, zone)
        by_name = {}
        for record in glue:
            by_name.setdefault((record.name, record.rtype), []).append(record)
        for (name, rtype), records in by_name.items():
            self.cache.put(name, rtype, records)
        addresses = [record.data for record in glue if record.rtype == w.TYPES["A"]]
        if addresses or depth > 4:
            return addresses
        # Sin glue, el nombre del servidor hay que resolverlo aparte, como cualquier otro.
        for record in referral:
            rcode, records, _ = self.resolve(record.data, w.TYPES["A"], depth + 1)
            if rcode == w.NOERROR:
                addresses += [r.data for r in records if r.rtype == w.TYPES["A"]]
        return addresses

    def remember_answer(self, reply):
        groups = {}
        for record in reply.answer:
            groups.setdefault((record.name, record.rtype), []).append(record)
        for (name, rtype), records in groups.items():
            self.cache.put(name, rtype, records)

    def remember_negative(self, name, rtype, reply):
        soa = [r for r in reply.authority if r.rtype == w.TYPES["SOA"]]
        if soa:
            self.cache.put(name, rtype, [], ttl=min(soa[0].ttl, soa[0].data[6]), rcode=reply.rcode)


def build_reply(resolver, data, peer):
    """Devuelve la respuesta completa a una pregunta, o None si no se puede ni contestar."""
    try:
        query = w.decode(data)
    except w.WireError:
        return None
    reply = w.Message(query.msg_id, w.FLAG_QR | w.FLAG_RA | (query.flags & w.FLAG_RD),
                      question=query.question[:1])
    if len(query.question) != 1:
        reply.flags |= w.FORMERR
        return reply
    name, rtype = query.question[0]
    if query.flags & w.FLAG_RD:
        rcode, records, authority = resolver.resolve(name, rtype)
    else:
        # Sin recursión, un resolver solo cuenta lo que ya tiene en cache.
        cached = resolver.cache.get(w.normalize(name), rtype)
        rcode, records, authority = (cached[1], cached[0], []) if cached else (w.NOERROR, [], [])
    reply.flags |= rcode
    reply.answer = records
    reply.authority = authority
    print(f"{peer[0]}:{peer[1]} id={query.msg_id} {name} {w.TYPE_NAMES.get(rtype, rtype)}"
          f" -> {w.RCODE_NAMES[rcode]} an={len(records)}", flush=True)
    return reply


def handle_udp(resolver, sock, data, peer):
    reply = build_reply(resolver, data, peer)
    if reply is None:
        return
    raw = w.encode(reply)
    if len(raw) > UDP_LIMIT:
        raw = w.encode(w.Message(reply.msg_id, reply.flags | w.FLAG_TC, question=reply.question))
    sock.sendto(raw, peer)


def handle_tcp(resolver, conn, peer):
    with conn:
        conn.settimeout(5)
        try:
            with conn.makefile("rb") as stream:
                head = stream.read(2)
                if len(head) != 2:
                    return
                data = stream.read(struct.unpack("!H", head)[0])
            reply = build_reply(resolver, data, peer)
            if reply is not None:
                raw = w.encode(reply)
                conn.sendall(struct.pack("!H", len(raw)) + raw)
        except OSError:
            return


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


def main():
    if len(sys.argv) != 3:
        sys.exit("uso: dns_resolver.py <dirección> <raíz>")
    resolver = Resolver(sys.argv[2])
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.bind((sys.argv[1], 53))
    threading.Thread(target=serve_tcp, args=(resolver, sys.argv[1]), daemon=True).start()
    print(f"resolver en {sys.argv[1]}, raíz {sys.argv[2]}", flush=True)
    while True:
        data, peer = sock.recvfrom(4096)
        threading.Thread(target=handle_udp, args=(resolver, sock, data, peer), daemon=True).start()


if __name__ == "__main__":
    main()