pounce/backend/app/services/yield_dns.py

"""
Yield DNS verification helpers.

Production-grade DNS checks for the Yield Connect flow:
- Option A (recommended): Nameserver delegation to our nameservers
- Option B (simpler): CNAME/ALIAS to a shared target
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional

import dns.resolver


@dataclass(frozen=True)
class YieldDNSCheckResult:
    verified: bool
    method: Optional[str]  # "nameserver" | "cname" | None
    actual_ns: list[str]
    cname_ok: bool
    error: Optional[str]


def _resolver() -> dns.resolver.Resolver:
    r = dns.resolver.Resolver()
    r.timeout = 3
    r.lifetime = 5
    return r


def _normalize_host(host: str) -> str:
    return host.rstrip(".").lower().strip()


def _resolve_ns(domain: str) -> list[str]:
    r = _resolver()
    answers = r.resolve(domain, "NS")
    # NS answers are RRset with .target
    return sorted({_normalize_host(str(rr.target)) for rr in answers})


def _resolve_cname(domain: str) -> list[str]:
    r = _resolver()
    answers = r.resolve(domain, "CNAME")
    return sorted({_normalize_host(str(rr.target)) for rr in answers})


def _resolve_a(host: str) -> list[str]:
    r = _resolver()
    answers = r.resolve(host, "A")
    return sorted({str(rr) for rr in answers})


def verify_yield_dns(domain: str, expected_nameservers: list[str], cname_target: str) -> YieldDNSCheckResult:
    """
    Verify that a domain is connected for Yield.

    We accept:
    - Nameserver delegation (NS contains all expected nameservers), OR
    - CNAME/ALIAS to `cname_target` (either CNAME matches, or A records match target A records)
    """
    domain = _normalize_host(domain)
    expected_ns = sorted({_normalize_host(ns) for ns in expected_nameservers if ns})
    target = _normalize_host(cname_target)

    if not domain:
        return YieldDNSCheckResult(
            verified=False,
            method=None,
            actual_ns=[],
            cname_ok=False,
            error="Domain is empty",
        )
    if not expected_ns and not target:
        return YieldDNSCheckResult(
            verified=False,
            method=None,
            actual_ns=[],
            cname_ok=False,
            error="Yield DNS is not configured on server",
        )

    # Option A: NS delegation
    try:
        actual_ns = _resolve_ns(domain)
        if expected_ns and set(expected_ns).issubset(set(actual_ns)):
            return YieldDNSCheckResult(
                verified=True,
                method="nameserver",
                actual_ns=actual_ns,
                cname_ok=False,
                error=None,
            )
    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
        actual_ns = []
    except Exception as e:
        return YieldDNSCheckResult(
            verified=False,
            method=None,
            actual_ns=[],
            cname_ok=False,
            error=str(e),
        )

    # Option B: CNAME / ALIAS
    if not target:
        return YieldDNSCheckResult(
            verified=False,
            method=None,
            actual_ns=actual_ns,
            cname_ok=False,
            error="Yield CNAME target is not configured on server",
        )

    # 1) Direct CNAME check (works for subdomain CNAME setups)
    try:
        cnames = _resolve_cname(domain)
        if any(c == target for c in cnames):
            return YieldDNSCheckResult(
                verified=True,
                method="cname",
                actual_ns=actual_ns,
                cname_ok=True,
                error=None,
            )
    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
        pass
    except Exception as e:
        return YieldDNSCheckResult(
            verified=False,
            method=None,
            actual_ns=actual_ns,
            cname_ok=False,
            error=str(e),
        )

    # 2) ALIAS/ANAME flattening: compare A records against target A records
    try:
        target_as = set(_resolve_a(target))
        domain_as = set(_resolve_a(domain))
        if target_as and domain_as and domain_as.issubset(target_as):
            return YieldDNSCheckResult(
                verified=True,
                method="cname",
                actual_ns=actual_ns,
                cname_ok=True,
                error=None,
            )
    except (dns.resolver.NXDOMAIN, dns.resolver.NoAnswer):
        pass
    except Exception as e:
        return YieldDNSCheckResult(
            verified=False,
            method=None,
            actual_ns=actual_ns,
            cname_ok=False,
            error=str(e),
        )

    return YieldDNSCheckResult(
        verified=False,
        method=None,
        actual_ns=actual_ns,
        cname_ok=False,
        error=None,
    )