"""Minimal logging DNS forwarder for the spens interceptor container.

Phase-2 DNS control: Docker's embedded DNS resolver (127.0.0.11) is
special-cased past the egress filtering of ``--internal`` networks, so an
agent on the isolated session network could otherwise resolve arbitrary
external names with no capture and no policy -- a metadata / low-bandwidth
exfiltration channel.  The runner therefore points the agent container's
resolv.conf at this forwarder (``docker run --dns <interceptor-ip>``), which

* enforces the **same domain allowlist** the mitmproxy addon applies to HTTP
  egress (``domain_rules`` in the interceptor config): a query is only
  forwarded to the upstream resolver if its name matches an allowed domain
  (suffix semantics) or one of the always-allowed internal names (the
  interceptor alias, localhost).  Any other name is **refused without being
  forwarded**, so an agent cannot smuggle data out by encoding it into the
  labels of an attacker-controlled domain and letting the recursive resolver
  leak it.  When *no* domain rules are configured there is no policy to
  enforce (same fail-open semantics as the HTTP path), so every name is
  forwarded,
* forwards allowed queries (UDP and TCP) unmodified to the container's
  normal upstream resolver (Docker's embedded DNS, via ``/etc/resolv.conf``),
* logs every query (client, name, query type, and whether it was allowed or
  blocked) to a JSONL file in the shared traces directory so name resolution
  shows up in the session audit trail,
* writes a readiness marker next to the mitmproxy addon's, so the runner and
  the agent entrypoint can block until DNS interception is actually up.

Stdlib only.  Runs as root inside the interceptor container (port 53 is
privileged); the agent never runs this code.
"""

from __future__ import annotations

import contextlib
import json
import os
import socket
import struct
import threading
from datetime import UTC, datetime
from pathlib import Path

try:
    # In the interceptor container both modules sit side by side in /app and
    # this file runs as a script, so the flat import is what resolves.
    from mitmproxy_helpers import domain_matches
except ImportError:
    # Imported as part of the installed spens package (tests, tooling), where
    # /app is not on the path.  build_interceptor_image copies both modules
    # into /app, so the dependency is present in either layout.
    from spens.data.mitmproxy_helpers import domain_matches

LISTEN_HOST = os.environ.get("SPENS_DNS_LISTEN", "0.0.0.0")
LISTEN_PORT = int(os.environ.get("SPENS_DNS_PORT", "53"))
LOG_PATH = os.environ.get("SPENS_DNS_LOG", "/app/traces/dns_log.jsonl")
RESOLV_CONF = os.environ.get("SPENS_DNS_RESOLV_CONF", "/etc/resolv.conf")
READY_MARKER_DIR = os.environ.get("SPENS_DNS_READY_DIR", "/root/.mitmproxy")
READY_MARKER_FILENAME = "spens_dns_ready"
UPSTREAM_TIMEOUT = float(os.environ.get("SPENS_DNS_TIMEOUT", "5"))
MAX_PACKET = 65535

# The interceptor config the mitmproxy addon reads; the DNS forwarder reads the
# same file so DNS resolution is filtered by the *same* ``domain_rules``
# allowlist that gates HTTP egress -- one policy, enforced on both channels.
POLICY_CONFIG_PATH = os.environ.get(
    "SPENS_DNS_POLICY_CONFIG", "/app/spens_interceptor_config.json"
)
# Internal names the agent must always be able to resolve to function even
# under a strict allowlist: the interceptor proxy alias (so it can reach
# mitmproxy at all) and loopback.  Comma-separated, overridable.
ALWAYS_ALLOW = [
    h.strip().lower()
    for h in os.environ.get(
        "SPENS_DNS_ALWAYS_ALLOW", "spens-interceptor,localhost"
    ).split(",")
    if h.strip()
]
RCODE_REFUSED = 0x0005

QTYPE_NAMES = {
    1: "A",
    2: "NS",
    5: "CNAME",
    6: "SOA",
    12: "PTR",
    15: "MX",
    16: "TXT",
    28: "AAAA",
    33: "SRV",
    35: "NAPTR",
    41: "OPT",
    43: "DS",
    46: "RRSIG",
    48: "DNSKEY",
    64: "SVCB",
    65: "HTTPS",
    252: "AXFR",
    255: "ANY",
}


def _timestamp() -> str:
    return datetime.now(UTC).isoformat().replace("+00:00", "Z")


def upstream_from_resolv_conf(path: str = "/etc/resolv.conf") -> tuple[str, int]:
    """Return the first ``nameserver`` from ``path`` (Docker's embedded DNS
    when attached to a user-defined network), falling back to 127.0.0.11."""
    try:
        with open(path, encoding="utf-8") as fh:
            for line in fh:
                parts = line.split()
                if len(parts) >= 2 and parts[0] == "nameserver":
                    return parts[1], 53
    except OSError:
        pass
    return "127.0.0.11", 53


def parse_questions(message: bytes) -> list[tuple[str, int]]:
    """Extract ``(name, qtype)`` pairs from a DNS message's question section.

    Tolerates malformed input: returns whatever was parsed so far (possibly
    nothing) instead of raising, so a hostile query can never crash the
    forwarder.  Names are decoded as latin-1 to preserve the raw bytes for
    the audit log.
    """
    questions: list[tuple[str, int]] = []
    try:
        qdcount = struct.unpack_from("!H", message, 4)[0]
        offset = 12
        for _ in range(qdcount):
            labels: list[str] = []
            while True:
                length = message[offset]
                offset += 1
                if length == 0:
                    break
                if length & 0xC0:
                    # compression pointers are not valid in a question section
                    return questions
                labels.append(message[offset : offset + length].decode("latin-1"))
                offset += length
            qtype = struct.unpack_from("!H", message, offset)[0]
            offset += 4
            questions.append((".".join(labels), qtype))
    except (IndexError, struct.error):
        pass
    return questions


def build_servfail(query: bytes) -> bytes:
    """Return a minimal SERVFAIL response for ``query`` (used when the
    upstream resolver cannot be reached)."""
    if len(query) < 12:
        return query
    flags = struct.unpack_from("!H", query, 2)[0]
    flags |= 0x8000  # QR: this is a response
    flags = (flags & 0xFFF0) | 0x0002  # RCODE = SERVFAIL
    return query[:2] + struct.pack("!H", flags) + query[4:]


def build_refused(query: bytes) -> bytes:
    """Return a minimal REFUSED response for ``query``.

    Sent for names blocked by the domain allowlist.  The query is **never**
    forwarded upstream, so no bytes of the (potentially data-bearing) name
    leave the interceptor -- this is what closes the recursive-DNS exfil
    channel.  RCODE REFUSED (rather than NXDOMAIN) signals a policy denial
    rather than a non-existent name.
    """
    if len(query) < 12:
        return query
    flags = struct.unpack_from("!H", query, 2)[0]
    flags |= 0x8000  # QR: this is a response
    flags = (flags & 0xFFF0) | RCODE_REFUSED
    return query[:2] + struct.pack("!H", flags) + query[4:]


def name_allowed(
    name: str,
    domain_rules: list[dict],
    always_allow: list[str] | None = None,
) -> bool:
    """Return whether ``name`` may be resolved under the domain allowlist.

    Delegates the pattern match to :func:`mitmproxy_helpers.domain_matches`,
    the same rule the mitmproxy addon applies to HTTP egress, so DNS and HTTP
    enforce one policy and cannot drift: a rule ``pattern`` matches the name
    itself or any of its subdomains, at a label boundary -- so
    ``anthropic.com`` matches ``api.anthropic.com`` but not
    ``anthropic.com.attacker.example`` nor ``evil-anthropic.com``.

    Fails **closed** when rules are configured but none match.  When no rules
    are configured at all there is no policy to enforce, so every name is
    allowed (mirrors ``is_allowed``/``host_allowed`` on the HTTP path).
    ``always_allow`` names (the interceptor alias, localhost) are permitted
    regardless, so the agent can always reach the proxy.
    """
    # An empty always_allow entry would strip to a bare wildcard and authorize
    # every name, so blank entries are skipped rather than matched.
    for allow in always_allow or []:
        if allow.strip() and domain_matches(name, allow):
            return True
    if not domain_rules:
        return True
    return any(
        domain_matches(name, str(rule.get("pattern", ""))) for rule in domain_rules
    )


def questions_allowed(
    questions: list[tuple[str, int]],
    domain_rules: list[dict],
    always_allow: list[str] | None = None,
) -> bool:
    """Allow a query only if **every** question is allowed (fail closed).

    A query with no parseable questions is refused when a policy is active:
    an unparseable/empty question section has no legitimate use here and
    must not be forwarded to the recursive resolver.
    """
    if not domain_rules:
        return True
    if not questions:
        return False
    return all(
        name_allowed(name, domain_rules, always_allow) for name, _ in questions
    )


def load_domain_rules(path: str = POLICY_CONFIG_PATH) -> list[dict]:
    """Load ``domain_rules`` from the interceptor config (best effort).

    A missing/malformed config yields an empty list (no DNS policy), matching
    the addon's behaviour when it cannot read its config.
    """
    try:
        with open(path, encoding="utf-8") as fh:
            data = json.load(fh)
    except (OSError, json.JSONDecodeError, ValueError):
        return []
    rules = data.get("domain_rules") if isinstance(data, dict) else None
    if not isinstance(rules, list):
        return []
    return [r for r in rules if isinstance(r, dict)]


def log_entry(
    questions: list[tuple[str, int]],
    client: str,
    proto: str,
    action: str = "forward",
) -> dict:
    """Build a JSONL-serializable audit entry for a DNS query.

    ``action`` records the policy decision: ``"forward"`` (sent upstream) or
    ``"blocked"`` (refused by the domain allowlist, never forwarded).
    """
    return {
        "type": "dns_query",
        "timestamp": _timestamp(),
        "client": client,
        "proto": proto,
        "action": action,
        "questions": [
            {"name": name, "type": QTYPE_NAMES.get(qtype, f"TYPE{qtype}")}
            for name, qtype in questions
        ],
    }


class DnsQueryLogger:
    """Append DNS query entries to a JSONL file, thread-safely."""

    def __init__(self, path: str | Path) -> None:
        self._path = Path(path)
        self._lock = threading.Lock()

    def log(self, entry: dict) -> None:
        try:
            with self._lock, open(self._path, "a", encoding="utf-8") as fh:
                fh.write(json.dumps(entry) + "\n")
        except OSError:
            # Logging must never take down name resolution for the agent.
            pass


def write_ready_marker(directory: Path) -> Path:
    """Write the DNS-forwarder readiness marker and return its path.

    Lives in the shared cert volume (mounted read-only at /certs in the
    agent), mirroring the addon's ``spens_addon_ready`` marker, so the runner
    and the agent entrypoint can block until DNS interception is up.
    """
    directory.mkdir(parents=True, exist_ok=True)
    marker = directory / READY_MARKER_FILENAME
    marker.write_text(_timestamp() + "\n", encoding="utf-8")
    return marker


def _read_tcp_message(sock: socket.socket) -> bytes | None:
    """Read one length-prefixed DNS-over-TCP message (None on EOF)."""
    header = b""
    while len(header) < 2:
        chunk = sock.recv(2 - len(header))
        if not chunk:
            return None
        header += chunk
    (length,) = struct.unpack("!H", header)
    data = b""
    while len(data) < length:
        chunk = sock.recv(length - len(data))
        if not chunk:
            return None
        data += chunk
    return data


class DnsForwarder:
    """UDP + TCP DNS forwarder that logs every query before forwarding it."""

    def __init__(
        self,
        upstream: tuple[str, int],
        logger: DnsQueryLogger,
        timeout: float = UPSTREAM_TIMEOUT,
        domain_rules: list[dict] | None = None,
        always_allow: list[str] | None = None,
    ) -> None:
        self._upstream = upstream
        self._logger = logger
        self._timeout = timeout
        self._domain_rules = domain_rules or []
        self._always_allow = always_allow if always_allow is not None else ALWAYS_ALLOW

    def _allowed(self, questions: list[tuple[str, int]]) -> bool:
        return questions_allowed(questions, self._domain_rules, self._always_allow)

    # -- UDP -------------------------------------------------------------

    def handle_udp(self, data: bytes, addr, server: socket.socket) -> None:
        questions = parse_questions(data)
        if not self._allowed(questions):
            # Blocked by the domain allowlist: log and refuse WITHOUT
            # forwarding, so the name never reaches the recursive resolver.
            self._logger.log(log_entry(questions, addr[0], "udp", action="blocked"))
            with contextlib.suppress(OSError):
                server.sendto(build_refused(data), addr)
            return
        self._logger.log(log_entry(questions, addr[0], "udp"))
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as fwd:
                fwd.settimeout(self._timeout)
                fwd.sendto(data, self._upstream)
                response, _ = fwd.recvfrom(MAX_PACKET)
        except OSError:
            response = build_servfail(data)
        with contextlib.suppress(OSError):
            server.sendto(response, addr)

    def serve_udp(self, sock: socket.socket) -> None:
        while True:
            data, addr = sock.recvfrom(MAX_PACKET)
            threading.Thread(
                target=self.handle_udp, args=(data, addr, sock), daemon=True
            ).start()

    # -- TCP -------------------------------------------------------------

    def handle_tcp(self, conn: socket.socket) -> None:
        try:
            peer = conn.getpeername()[0]
        except OSError:
            peer = "?"
        try:
            conn.settimeout(self._timeout * 4)
            with socket.create_connection(self._upstream, timeout=self._timeout) as fwd:
                while True:
                    query = _read_tcp_message(conn)
                    if query is None:
                        return
                    questions = parse_questions(query)
                    if not self._allowed(questions):
                        self._logger.log(
                            log_entry(questions, peer, "tcp", action="blocked")
                        )
                        refused = build_refused(query)
                        conn.sendall(struct.pack("!H", len(refused)) + refused)
                        continue
                    self._logger.log(log_entry(questions, peer, "tcp"))
                    fwd.sendall(struct.pack("!H", len(query)) + query)
                    response = _read_tcp_message(fwd)
                    if response is None:
                        return
                    conn.sendall(struct.pack("!H", len(response)) + response)
        except OSError:
            pass
        finally:
            conn.close()

    def serve_tcp(self, sock: socket.socket) -> None:
        while True:
            conn, _ = sock.accept()
            threading.Thread(target=self.handle_tcp, args=(conn,), daemon=True).start()


def main() -> None:
    log_dir = os.path.dirname(os.path.abspath(LOG_PATH))
    Path(log_dir).mkdir(parents=True, exist_ok=True)
    upstream = upstream_from_resolv_conf(RESOLV_CONF)
    domain_rules = load_domain_rules(POLICY_CONFIG_PATH)
    forwarder = DnsForwarder(
        upstream,
        DnsQueryLogger(LOG_PATH),
        domain_rules=domain_rules,
        always_allow=ALWAYS_ALLOW,
    )
    if domain_rules:
        print(
            f"[spens-dns] domain allowlist active: {len(domain_rules)} rule(s); "
            f"always-allow {ALWAYS_ALLOW}; non-matching names are REFUSED "
            "(not forwarded)",
            flush=True,
        )
    else:
        print(
            "[spens-dns] no domain_rules configured -- DNS is not filtered "
            "(fail-open, matching HTTP egress policy)",
            flush=True,
        )

    # Bind both listeners BEFORE writing the readiness marker, so anyone
    # gating on the marker knows port 53 is actually serving.
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    udp_sock.bind((LISTEN_HOST, LISTEN_PORT))
    tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp_sock.bind((LISTEN_HOST, LISTEN_PORT))
    tcp_sock.listen(16)

    print(
        f"[spens-dns] forwarding {LISTEN_HOST}:{LISTEN_PORT} -> "
        f"{upstream[0]}:{upstream[1]}, logging to {LOG_PATH}",
        flush=True,
    )

    threading.Thread(target=forwarder.serve_tcp, args=(tcp_sock,), daemon=True).start()
    write_ready_marker(Path(READY_MARKER_DIR))
    # Serve UDP on the main thread: if it dies, the process exits loudly
    # instead of silently keeping a dead DNS server around.
    forwarder.serve_udp(udp_sock)


if __name__ == "__main__":
    main()
