app/validator/doh_resolver.py

"""DNS/MX worker over DNS-over-HTTPS (DoH) — free public resolvers.

Architecture spec §3.2 / §7.3:

- MX resolution runs on Google Public DNS and Cloudflare DNS over HTTPS.
  Both are free, unlimited and need no authentication; HTTPS/443 egress is
  never blocked (unlike UDP/53 on locked-down PaaS hosts).
- The SAME query goes to BOTH resolvers. Divergent MX sets flag the domain as
  `dns_inconsistent` — a deterministic risk signal (regional DNS hijack /
  misconfiguration), requiring no inference.
- Results are cached at the domain level: 24h TTL for MX answers, 1h for
  negative (null-MX / NXDOMAIN) answers per the spec.
- Punycode normalization happens here (§7.1) so Unicode domains resolve.

The legacy dnspython path in dns_check.py remains the fallback: any DoH
failure (network, JSON, schema) degrades to it transparently. The DoH layer
can also be disabled entirely with DOH_ENABLED=false.
"""
from __future__ import annotations

import asyncio
import json
from typing import Any, Optional

import httpx

from ..config import settings
from .i18n import to_punycode

GOOGLE_DOH = "https://dns.google/resolve"
CLOUDFLARE_DOH = "https://cloudflare-dns.com/dns-query"

MX_TTL_SECONDS = 24 * 3600          # spec: 24h for MX records
NEGATIVE_TTL_SECONDS = 1 * 3600     # spec: 1h for null-MX negatives

_TIMEOUT_LABEL = "timeout"


async def _doh_query(
    client: httpx.AsyncClient,
    endpoint: str,
    name: str,
    rtype: str,
    extra_params: dict | None = None,
    accept: str = "application/dns-json",
) -> dict:
    """One DoH JSON query (RFC 8484-style JSON APIs used by both providers)."""
    params = {"name": name, "type": rtype}
    if extra_params:
        params.update(extra_params)
    resp = await client.get(endpoint, params=params, headers={"accept": accept})
    resp.raise_for_status()
    data = resp.json()
    if not isinstance(data, dict):
        raise ValueError("doh: non-object response")
    return data


def _parse_mx(data: dict) -> list[dict]:
    """Extract MX records (sorted by priority) from a DoH answer."""
    records: list[dict] = []
    for ans in data.get("Answer") or []:
        # RFC 1035 type 15 = MX
        if ans.get("type") != 15:
            continue
        parts = str(ans.get("data", "")).split()
        if len(parts) >= 2:
            try:
                priority = int(parts[0])
            except ValueError:
                priority = 0
            host = " ".join(parts[1:]).rstrip(".").strip().lower()
            if host:
                records.append({"host": host, "priority": priority})
    records.sort(key=lambda r: (r["priority"], r["host"]))
    return records


def _parse_a_like(data: dict, rtypes: set[int]) -> bool:
    """True when the answer holds at least one A (1) / AAAA (28) record."""
    for ans in data.get("Answer") or []:
        if ans.get("type") in rtypes:
            return True
    return False


def _mx_sets_match(a: list[dict], b: list[dict]) -> bool:
    """Set-equality on (priority, host) — order-insensitive, tolerant."""
    norm = lambda rs: {(r["host"], r["priority"]) for r in rs}
    return norm(a) == norm(b)


async def resolve_mx_dual(
    domain: str, timeout: Optional[float] = None
) -> dict[str, Any]:
    """Query BOTH public resolvers for MX and compare.

    Returns:
        {
          "status": "ok" | "nxdomain" | "no_mail" | "error",
          "mx_records": [...],          # primary (Google) answer
          "consistent": bool | None,    # None when the second resolver failed
          "a_fallback": bool,
          "error": str | None,
        }
    """
    domain = to_punycode((domain or "").strip().lower().rstrip("."))
    timeout = timeout if timeout is not None else settings.doh_timeout

    async def _query_both() -> tuple[dict | None, dict | None, Exception | None]:
        async with httpx.AsyncClient(timeout=timeout) as client:
            google_task = _doh_query(client, GOOGLE_DOH, domain, "MX")
            cf_task = _doh_query(
                client, CLOUDFLARE_DOH, domain, "MX", extra_params={"ct": "application/dns-json"}
            )
            g_res = await asyncio.gather(google_task, cf_task, return_exceptions=True)
        g, c = g_res
        g_err = g if isinstance(g, Exception) else None
        c_err = c if isinstance(c, Exception) else None
        g_data = g if isinstance(g, dict) else None
        c_data = c if isinstance(c, dict) else None
        # If the PRIMARY (Google) failed but Cloudflare answered, swap so the
        # pipeline still gets an authoritative answer from a public resolver.
        if g_data is None and c_data is not None:
            g_data, c_data = c_data, None
        return g_data, c_data, g_err or c_err

    try:
        g_data, c_data, err = await _query_both()
    except Exception as exc:  # httpx transport failure etc.
        return {"status": "error", "mx_records": [], "consistent": None,
                "a_fallback": False, "error": f"doh transport: {exc}"}

    if g_data is None:
        return {"status": "error", "mx_records": [], "consistent": None,
                "a_fallback": False, "error": f"doh unavailable: {err}"}

    g_status = g_data.get("Status", 0)
    if g_status == 3:  # NXDOMAIN
        return {"status": "nxdomain", "mx_records": [], "consistent": None,
                "a_fallback": False, "error": "Domain does not exist (NXDOMAIN)"}
    if g_status not in (0, 3):
        return {"status": "error", "mx_records": [], "consistent": None,
                "a_fallback": False, "error": f"doh resolver status {g_status}"}

    mx = _parse_mx(g_data)
    if mx:
        consistent = (
            _mx_sets_match(mx, _parse_mx(c_data))
            if c_data is not None and c_data.get("Status", 0) == 0
            else None
        )
        return {"status": "ok", "mx_records": mx, "consistent": consistent,
                "a_fallback": False, "error": None}

    # No MX → A/AAAA fallback per RFC 5321 §5.1
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            a_task = _doh_query(client, GOOGLE_DOH, domain, "A")
            aaaa_task = _doh_query(client, GOOGLE_DOH, domain, "AAAA")
            a_res = await asyncio.gather(a_task, aaaa_task, return_exceptions=True)
        has_addr = any(
            isinstance(d, dict) and _parse_a_like(d, {1, 28}) for d in a_res
        )
    except Exception:
        has_addr = False

    if has_addr:
        return {"status": "ok", "a_fallback": True,
                "mx_records": [{"host": domain, "priority": 0}],
                "consistent": None, "error": None}
    return {"status": "no_mail", "mx_records": [], "consistent": None,
            "a_fallback": False, "error": "Domain has no MX, A, or AAAA records"}


# ---- domain-level TTL cache (shared-cache backed, memory tier always works) --

async def _cache_get(key: str) -> dict | None:
    try:
        from ..utils.cache import get_cache
        raw = await get_cache(key)
        if not raw:
            return None
        data = json.loads(raw)
        return data if isinstance(data, dict) else None
    except Exception:
        return None  # cache outage → plain re-resolve


async def _cache_set(key: str, payload: dict, ttl: int) -> None:
    try:
        from ..utils.cache import set_cache
        await set_cache(key, json.dumps(payload), ttl)
    except Exception:
        pass  # a cache write failure must never break resolution


def _mx_cache_key(domain: str) -> str:
    return f"doh:mx:{to_punycode(domain.strip().lower())}"


async def resolve_mx_dual_cached(domain: str, timeout: Optional[float] = None) -> dict:
    """`resolve_mx_dual` with a domain-level TTL cache.

    Definitive answers cache for 24h; negative answers (NXDOMAIN / no mail)
    for 1h per the spec. Errors are never cached.
    """
    domain = (domain or "").strip().lower()
    key = _mx_cache_key(domain)
    cached = await _cache_get(key)
    if cached is not None:
        cached = {**cached, "cache_hit": True}
        return cached

    result = await resolve_mx_dual(domain, timeout)
    if result.get("status") == "ok":
        await _cache_set(key, {**result, "cache_hit": False}, MX_TTL_SECONDS)
    elif result.get("status") in {"nxdomain", "no_mail"}:
        await _cache_set(key, {**result, "cache_hit": False}, NEGATIVE_TTL_SECONDS)
    return result