#!/usr/bin/env python3
"""A deliberately tiny DNS responder that answers ONE name and forwards the rest.

WHY THIS EXISTS. FIFA 17's ProtoSSL verifies the roster certificate by dNSName
only, so the client must reach the roster as `winter15.gosredirector.ea.com`
rather than by IP -- an IP-addressed roster is rejected and the player gets
"An error occurred downloading the FUT Squad Update". The usual fix is a hosts
entry on the client, which on this operator's machine previously took the whole
machine's internet down.

So instead of editing the client's hosts file, the client points ONE name at
this resolver through a Windows NRPT rule. Everything else on the client keeps
using its normal DNS and is never routed here at all.

SAFETY PROPERTY THAT MATTERS: even for the queries that do arrive, an unknown
name is FORWARDED upstream verbatim rather than refused or NXDOMAINed. So the
worst case if the NRPT scope is ever widened by accident is added latency, not a
broken resolver -- this cannot repeat the earlier outage.

Answers A records only; every other qtype for the pinned name is forwarded too,
because inventing an answer for a type we do not model is how a resolver starts
lying.
"""

import os
import socket
import socketserver
import struct
import sys

PINNED_NAME = os.environ.get("SCOPED_DNS_NAME", "winter15.gosredirector.ea.com").rstrip(".").lower()
PINNED_ADDR = os.environ.get("SCOPED_DNS_ADDR", "10.10.0.120")
BIND_ADDR = os.environ.get("SCOPED_DNS_BIND", "0.0.0.0")
BIND_PORT = int(os.environ.get("SCOPED_DNS_PORT", "53"))
TTL = 60

QTYPE_A = 1
QCLASS_IN = 1


def upstream_servers() -> list[str]:
    """Real resolvers, read from resolv.conf, minus ourselves."""
    out = []
    try:
        with open("/etc/resolv.conf") as fh:
            for line in fh:
                parts = line.split()
                if len(parts) >= 2 and parts[0] == "nameserver" and parts[1] != PINNED_ADDR:
                    out.append(parts[1])
    except OSError:
        pass
    return out or ["1.1.1.1", "8.8.8.8"]


UPSTREAM = upstream_servers()


def parse_question(data: bytes):
    """Return (qname, qtype, end_offset) or None if the packet is not parseable."""
    if len(data) < 12:
        return None
    qdcount = struct.unpack("!H", data[4:6])[0]
    if qdcount < 1:
        return None
    labels, off = [], 12
    while off < len(data):
        ln = data[off]
        if ln == 0:
            off += 1
            break
        # A pointer in the QUESTION section is malformed; forward and let the
        # upstream decide rather than guessing.
        if ln & 0xC0:
            return None
        off += 1
        labels.append(data[off:off + ln].decode("ascii", "replace"))
        off += ln
    if off + 4 > len(data):
        return None
    qtype, _qclass = struct.unpack("!HH", data[off:off + 4])
    return ".".join(labels).lower(), qtype, off + 4


def build_answer(query: bytes, qend: int) -> bytes:
    """Echo the question and append one A record for the pinned address."""
    txid = query[:2]
    # QR=1, RD copied from the query, RA=1.
    rd = query[2] & 0x01
    flags = struct.pack("!H", 0x8000 | (rd << 8) | 0x0080)
    counts = struct.pack("!HHHH", 1, 1, 0, 0)
    question = query[12:qend]
    rr = (
        b"\xc0\x0c"                                   # name -> pointer to the question
        + struct.pack("!HHIH", QTYPE_A, QCLASS_IN, TTL, 4)
        + socket.inet_aton(PINNED_ADDR)
    )
    return txid + flags + counts + question + rr


def forward(query: bytes) -> bytes | None:
    for server in UPSTREAM:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                s.settimeout(3.0)
                s.sendto(query, (server, 53))
                return s.recv(4096)
        except OSError:
            continue
    return None


class Handler(socketserver.BaseRequestHandler):
    def handle(self):
        data, sock = self.request
        parsed = parse_question(data)
        if parsed:
            qname, qtype, qend = parsed
            if qname == PINNED_NAME and qtype == QTYPE_A:
                sock.sendto(build_answer(data, qend), self.client_address)
                print(f"ANSWER {qname} A -> {PINNED_ADDR}  (from {self.client_address[0]})", flush=True)
                return
        reply = forward(data)
        if reply:
            sock.sendto(reply, self.client_address)
            if parsed:
                print(f"forward {parsed[0]} qtype={parsed[1]}  (from {self.client_address[0]})", flush=True)


class Server(socketserver.ThreadingUDPServer):
    allow_reuse_address = True
    daemon_threads = True


if __name__ == "__main__":
    print(
        f"scoped-dns: pinning {PINNED_NAME} -> {PINNED_ADDR}; "
        f"forwarding everything else to {', '.join(UPSTREAM)}",
        flush=True,
    )
    try:
        with Server((BIND_ADDR, BIND_PORT), Handler) as srv:
            print(f"scoped-dns: LISTENING on {BIND_ADDR}:{BIND_PORT}", flush=True)
            srv.serve_forever()
    except PermissionError:
        sys.exit(f"scoped-dns: cannot bind {BIND_ADDR}:{BIND_PORT} (needs root)")