webmail-client/backend/app/accounts/discovery.py

"""Autodiscovery resolver (SPEC §6.2).

``discover(email)`` runs the chain — presets, Thunderbird autoconfig, ISPDB, DNS
SRV (RFC 6186), then a live probe — returning on the first success but recording
every step's outcome for wizard UX. All HTTP is SSRF-guarded (§10.6).
"""

from __future__ import annotations

import asyncio
import ssl
from pathlib import Path

import dns.asyncresolver
import dns.exception
import httpx
import yaml
from defusedxml import ElementTree as SafeET

from app.accounts.ssrf import SsrfError, assert_host_is_public, guarded_get
from app.accounts.types import (
    CandidateConfig,
    DiscoveryResult,
    DiscoveryStep,
    HelpBlock,
    ServerConfig,
)

_PROBE_TIMEOUT = 5.0

_presets_cache: dict[str, CandidateConfig] | None = None


# --- presets -----------------------------------------------------------------


def _strategy_from_auth(auth: str) -> tuple[str, ...]:
    """Default credential flow when a preset doesn't declare one (SPEC §6.2.1)."""
    return ("oauth2",) if auth == "oauth2" else ("password",)


def _help_block(raw: object) -> HelpBlock | None:
    if not isinstance(raw, dict):
        return None
    steps = raw.get("steps", [])
    return HelpBlock(
        help_id=str(raw["help_id"]),
        deep_link=str(raw["deep_link"]),
        app_name_hint=str(raw.get("app_name_hint", "")),
        steps=tuple(str(s) for s in steps) if isinstance(steps, list) else (),
        docs_url=str(raw.get("docs_url", "")),
        setup_link=str(raw.get("setup_link", "")),
        button_after_step=int(raw.get("button_after_step", 0)),
    )


def _load_presets() -> dict[str, CandidateConfig]:
    global _presets_cache
    if _presets_cache is None:
        raw = yaml.safe_load((Path(__file__).parent / "presets.yaml").read_text())
        mapping: dict[str, CandidateConfig] = {}
        for prov in raw["providers"].values():
            auth = prov.get("auth", "password")
            imap = ServerConfig(
                prov["imap"]["host"], prov["imap"]["port"], prov["imap"]["tls"], auth
            )
            smtp = ServerConfig(
                prov["smtp"]["host"], prov["smtp"]["port"], prov["smtp"]["tls"], auth
            )
            strategy = tuple(prov.get("auth_strategy", [])) or _strategy_from_auth(auth)
            candidate = CandidateConfig(
                imap,
                smtp,
                "preset",
                prov.get("oauth_provider"),
                strategy,
                _help_block(prov.get("help")),
            )
            for domain in prov["domains"]:
                mapping[domain.lower()] = candidate
        _presets_cache = mapping
    return _presets_cache


def domain_of(email: str) -> str:
    return email.rsplit("@", 1)[-1].strip().lower()


# --- Thunderbird autoconfig / ISPDB XML --------------------------------------


def _socket_type(text: str | None) -> str:
    return {"ssl": "ssl", "starttls": "starttls", "plain": "plain"}.get(
        (text or "").strip().lower(), "ssl"
    )


def parse_autoconfig_xml(xml_text: str) -> CandidateConfig | None:
    """Parse Thunderbird ``clientConfig`` XML into a CandidateConfig (or None)."""
    try:
        root = SafeET.fromstring(xml_text)
    except Exception:
        return None
    provider = root.find("emailProvider")
    if provider is None:
        return None

    def _server(kind: str, server_type: str) -> ServerConfig | None:
        for node in provider.findall(kind):
            if node.get("type") == server_type:
                host = node.findtext("hostname")
                port = node.findtext("port")
                if host and port and port.isdigit():
                    auth = node.findtext("authentication") or ""
                    is_oauth = "oauth" in auth.lower()
                    return ServerConfig(
                        host.strip(),
                        int(port),
                        _socket_type(node.findtext("socketType")),
                        "oauth2" if is_oauth else "password",
                    )
        return None

    imap = _server("incomingServer", "imap")
    smtp = _server("outgoingServer", "smtp")
    if imap is None or smtp is None:
        return None
    return CandidateConfig(
        imap, smtp, "autoconfig", auth_strategy=_strategy_from_auth(imap.auth)
    )


# --- DNS SRV (RFC 6186) ------------------------------------------------------


async def _srv_lookup(name: str, tls: str) -> ServerConfig | None:
    resolver = dns.asyncresolver.Resolver()
    resolver.lifetime = 5.0
    try:
        answer = await resolver.resolve(name, "SRV")
    except (dns.exception.DNSException, OSError):
        return None
    for record in sorted(answer, key=lambda r: (r.priority, r.weight)):
        target = str(record.target).rstrip(".")
        if target and target != ".":  # "." target => service explicitly unavailable
            return ServerConfig(target, record.port, tls)
    return None


async def _discover_srv(domain: str) -> CandidateConfig | None:
    imap = await _srv_lookup(f"_imaps._tcp.{domain}", "ssl") or await _srv_lookup(
        f"_imap._tcp.{domain}", "starttls"
    )
    smtp = await _srv_lookup(f"_submissions._tcp.{domain}", "ssl") or await _srv_lookup(
        f"_submission._tcp.{domain}", "starttls"
    )
    if imap and smtp:
        return CandidateConfig(imap, smtp, "srv")
    return None


# --- live probe --------------------------------------------------------------


async def _probe_one(host: str, port: int, tls: str, kind: str) -> ServerConfig | None:
    try:
        await assert_host_is_public(host)  # SSRF: never probe internal hosts
    except SsrfError:
        return None
    try:
        if tls == "ssl":
            ctx = ssl.create_default_context()
            conn = asyncio.open_connection(host, port, ssl=ctx)
        else:
            conn = asyncio.open_connection(host, port)
        reader, writer = await asyncio.wait_for(conn, _PROBE_TIMEOUT)
        greeting = await asyncio.wait_for(reader.readline(), _PROBE_TIMEOUT)
        writer.close()
    except (TimeoutError, OSError, ssl.SSLError):
        return None
    text = greeting.decode("ascii", "ignore")
    if kind == "imap" and text.startswith("* OK"):
        return ServerConfig(host, port, tls)
    if kind == "smtp" and text.startswith("220"):
        return ServerConfig(host, port, tls)
    return None


async def _first_working(candidates: list[tuple[str, int, str]], kind: str) -> ServerConfig | None:
    results = await asyncio.gather(
        *(_probe_one(host, port, tls, kind) for host, port, tls in candidates)
    )
    for result in results:  # candidates are in priority order
        if result is not None:
            return result
    return None


async def _discover_probe(domain: str) -> CandidateConfig | None:
    imap_hosts = [
        (f"imap.{domain}", 993, "ssl"),
        (f"mail.{domain}", 993, "ssl"),
        (f"imap.{domain}", 143, "starttls"),
        (f"mail.{domain}", 143, "starttls"),
    ]
    smtp_hosts = [
        (f"smtp.{domain}", 465, "ssl"),
        (f"mail.{domain}", 465, "ssl"),
        (f"smtp.{domain}", 587, "starttls"),
        (f"mail.{domain}", 587, "starttls"),
    ]
    imap, smtp = await asyncio.gather(
        _first_working(imap_hosts, "imap"), _first_working(smtp_hosts, "smtp")
    )
    if imap and smtp:
        return CandidateConfig(imap, smtp, "probe")
    return None


# --- orchestration -----------------------------------------------------------


async def _try_http_config(url: str, source: str) -> CandidateConfig | None:
    xml_text = await guarded_get(url)
    if xml_text is None:
        return None
    candidate = parse_autoconfig_xml(xml_text)
    if candidate is None:
        return None
    return CandidateConfig(
        candidate.imap,
        candidate.smtp,
        source,
        candidate.oauth_provider,
        candidate.auth_strategy,
        candidate.help,
    )


async def discover(email: str) -> DiscoveryResult:
    """Run the autodiscovery chain for ``email`` (SPEC §6.2)."""
    result = DiscoveryResult()
    domain = domain_of(email)

    # 1. Presets.
    presets = _load_presets()
    if domain in presets:
        result.candidates.append(presets[domain])
        result.steps.append(DiscoveryStep("preset", "found", domain))
        return result
    result.steps.append(DiscoveryStep("preset", "not_found", domain))

    # 2. Autoconfig (autoconfig.<domain>, then <domain>/.well-known).
    autoconfig_urls = [
        f"https://autoconfig.{domain}/mail/config-v1.1.xml?emailaddress={email}",
        f"https://{domain}/.well-known/autoconfig/mail/config-v1.1.xml?emailaddress={email}",
    ]
    for url in autoconfig_urls:
        try:
            candidate = await _try_http_config(url, "autoconfig")
        except (SsrfError, httpx.HTTPError) as exc:
            result.steps.append(DiscoveryStep("autoconfig", "error", str(exc)[:200]))
            continue
        if candidate is not None:
            result.candidates.append(candidate)
            result.steps.append(DiscoveryStep("autoconfig", "found", url))
            return result
    result.steps.append(DiscoveryStep("autoconfig", "not_found", domain))

    # 3. ISPDB (Thunderbird's database).
    ispdb_url = f"https://autoconfig.thunderbird.net/v1.1/{domain}"
    try:
        candidate = await _try_http_config(ispdb_url, "ispdb")
        if candidate is not None:
            result.candidates.append(candidate)
            result.steps.append(DiscoveryStep("ispdb", "found", ispdb_url))
            return result
        result.steps.append(DiscoveryStep("ispdb", "not_found", domain))
    except (SsrfError, httpx.HTTPError) as exc:
        result.steps.append(DiscoveryStep("ispdb", "error", str(exc)[:200]))

    # 4. DNS SRV (RFC 6186).
    srv = await _discover_srv(domain)
    if srv is not None:
        result.candidates.append(srv)
        result.steps.append(DiscoveryStep("srv", "found", domain))
        return result
    result.steps.append(DiscoveryStep("srv", "not_found", domain))

    # 5. Live probe.
    probe = await _discover_probe(domain)
    if probe is not None:
        result.candidates.append(probe)
        result.steps.append(DiscoveryStep("probe", "found", domain))
        return result
    result.steps.append(DiscoveryStep("probe", "not_found", domain))

    return result