mail-address-semantics/mas_core/discovery.py

"""Exact-host MAS statute discovery.

`discover(host)` fetches https://HOST/.well-known/mas and returns either a
`Discovered` (validated statute) or a typed `DiscoveryFailure`. It never
walks parent domains. Redirects are followed only within the same origin
(scheme, host, port); cross-host redirects and scheme downgrades fail.

Discovery failure never invalidates an email address; it only means MAS
namespace semantics are unconfirmed.
"""

from __future__ import annotations

import http.client
import json
import ssl
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Optional, Union
from urllib.parse import urljoin, urlsplit

from .statute import (
    SUPPORTED_VERSIONS,
    InheritanceCycle,
    Statute,
    StatuteInvalid,
    build_statute,
    canonical_host,
    resolve_inheritance,
)

WELL_KNOWN_PATH = "/.well-known/mas"
MAX_BYTES = 256 * 1024
MAX_REDIRECTS = 5


class FailureKind(str, Enum):
    InvalidHost = "InvalidHost"
    NotFound = "NotFound"
    HttpFailure = "HttpFailure"
    TransportFailure = "TransportFailure"
    InvalidContentType = "InvalidContentType"
    CrossHostRedirect = "CrossHostRedirect"
    SchemeDowngrade = "SchemeDowngrade"
    TooManyRedirects = "TooManyRedirects"
    NotJson = "NotJson"
    UnsupportedVersion = "UnsupportedVersion"
    SchemaInvalid = "SchemaInvalid"
    SemanticInvalid = "SemanticInvalid"
    NamespaceMismatch = "NamespaceMismatch"
    InheritanceCycle = "InheritanceCycle"
    InheritanceFailure = "InheritanceFailure"


@dataclass(frozen=True)
class DiscoveryFailure:
    host: str
    url: str
    kind: FailureKind
    detail: str
    cached: bool = False
    ok = False

    def to_dict(self) -> dict:
        return {"ok": False, "host": self.host, "url": self.url, "failure": self.kind.value,
                "detail": self.detail, "cached": self.cached}


@dataclass(frozen=True)
class Discovered:
    host: str
    url: str
    statute: Statute
    cached: bool = False
    ok = True

    def to_dict(self) -> dict:
        return {"ok": True, "host": self.host, "url": self.url, "cached": self.cached,
                "statute": self.statute.to_dict()}


DiscoveryResult = Union[Discovered, DiscoveryFailure]


@dataclass(frozen=True)
class Response:
    status: int
    headers: dict  # lower-cased header names
    body: bytes


Transport = Callable[[str, float], Response]


def http_transport(url: str, timeout: float) -> Response:
    """One GET, no redirect following, bounded body. Raises OSError on failure."""
    parts = urlsplit(url)
    if parts.scheme == "https":
        conn = http.client.HTTPSConnection(parts.hostname, parts.port, timeout=timeout,
                                           context=ssl.create_default_context())
    else:
        conn = http.client.HTTPConnection(parts.hostname, parts.port, timeout=timeout)
    try:
        conn.request("GET", parts.path or "/", headers={"Accept": "application/json",
                                                         "User-Agent": "mas-core/0.1"})
        resp = conn.getresponse()
        body = resp.read(MAX_BYTES + 1)
        return Response(resp.status, {k.lower(): v for k, v in resp.getheaders()}, body)
    finally:
        conn.close()


class _Fail(Exception):
    def __init__(self, kind: FailureKind, detail: str, url: str):
        super().__init__(detail)
        self.kind, self.detail, self.url = kind, detail, url


def _origin(url: str):
    p = urlsplit(url)
    default = {"https": 443, "http": 80}.get(p.scheme)
    return (p.scheme, (p.hostname or "").lower(), p.port or default)


def _fetch_json(url: str, transport: Transport, timeout: float) -> tuple[str, object]:
    """GET with same-origin redirects only; returns (final_url, decoded JSON)."""
    start = _origin(url)
    for _ in range(MAX_REDIRECTS + 1):
        try:
            resp = transport(url, timeout)
        except (OSError, http.client.HTTPException) as e:
            raise _Fail(FailureKind.TransportFailure, f"{type(e).__name__}: {e}", url)
        if resp.status in (301, 302, 303, 307, 308):
            loc = resp.headers.get("location")
            if not loc:
                raise _Fail(FailureKind.HttpFailure, f"{resp.status} without Location", url)
            nxt = urljoin(url, loc)
            o = _origin(nxt)
            if o[0] != start[0]:
                kind = FailureKind.SchemeDowngrade if o[0] == "http" else FailureKind.CrossHostRedirect
                raise _Fail(kind, f"redirect to {nxt} changes scheme", url)
            if o != start:
                raise _Fail(FailureKind.CrossHostRedirect,
                            f"redirect to {nxt} would transfer MAS authority", url)
            url = nxt
            continue
        if resp.status == 404 or resp.status == 410:
            raise _Fail(FailureKind.NotFound, f"HTTP {resp.status}", url)
        if not 200 <= resp.status < 300:
            raise _Fail(FailureKind.HttpFailure, f"HTTP {resp.status}", url)
        ctype = resp.headers.get("content-type", "").split(";")[0].strip().lower()
        if not (ctype == "application/json" or (ctype.startswith("application/") and ctype.endswith("+json"))):
            raise _Fail(FailureKind.InvalidContentType, f"content-type {ctype or '(none)'!r}", url)
        if len(resp.body) > MAX_BYTES:
            raise _Fail(FailureKind.HttpFailure, f"statute exceeds {MAX_BYTES} bytes", url)
        try:
            return url, json.loads(resp.body.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as e:
            raise _Fail(FailureKind.NotJson, str(e), url)
    raise _Fail(FailureKind.TooManyRedirects, f"more than {MAX_REDIRECTS} redirects", url)


def _validate(doc, host: str, url: str) -> Statute:
    if isinstance(doc, dict) and "mas_version" in doc and doc["mas_version"] not in SUPPORTED_VERSIONS:
        raise _Fail(FailureKind.UnsupportedVersion, f"mas_version {doc['mas_version']!r}", url)
    try:
        return build_statute(doc, host, statute_url=url)
    except StatuteInvalid as e:
        if "namespace-mismatch" in e.codes:
            raise _Fail(FailureKind.NamespaceMismatch, str(e), url)
        if "inheritance-cycle" in e.codes:
            raise _Fail(FailureKind.InheritanceCycle, str(e), url)
        kind = FailureKind.SchemaInvalid if e.layer == "schema" else FailureKind.SemanticInvalid
        raise _Fail(kind, str(e), url)


def statute_url(host: str, *, scheme: str = "https", port: Optional[int] = None) -> str:
    netloc = host if port is None else f"{host}:{port}"
    return f"{scheme}://{netloc}{WELL_KNOWN_PATH}"


@dataclass
class DiscoveryCache:
    """Minimal cache. Successes and failures are stored and expire separately;
    a cached failure is returned as the same typed failure, marked cached."""

    success_ttl: float = 3600.0
    failure_ttl: float = 60.0
    clock: Callable[[], float] = time.monotonic
    _ok: dict = field(default_factory=dict)
    _fail: dict = field(default_factory=dict)

    def get(self, key) -> Optional[DiscoveryResult]:
        for store in (self._ok, self._fail):
            hit = store.get(key)
            if hit and hit[0] > self.clock():
                r = hit[1]
                return type(r)(**{**r.__dict__, "cached": True})
            store.pop(key, None)
        return None

    def put(self, key, result: DiscoveryResult) -> None:
        if result.ok:
            self._fail.pop(key, None)
            self._ok[key] = (self.clock() + self.success_ttl, result)
        else:
            self._fail[key] = (self.clock() + self.failure_ttl, result)


def discover(host: str, *, transport: Transport = http_transport, timeout: float = 5.0,
             scheme: str = "https", port: Optional[int] = None,
             resolve_inherits: bool = True,
             cache: Optional[DiscoveryCache] = None) -> DiscoveryResult:
    """Discover the statute for exactly `host`. Never consults parent domains.

    `scheme="http"` and `port` exist only for local fixture servers in tests;
    production discovery is HTTPS on the default port.
    """
    if scheme not in ("https", "http"):
        raise ValueError("scheme must be https (or http for local test fixtures)")
    raw_host = host
    host = canonical_host(host)
    url = statute_url(host, scheme=scheme, port=port)
    if not host or any(c in host for c in "/@:?# ") :
        return DiscoveryFailure(raw_host, url, FailureKind.InvalidHost, "not a hostname")
    key = (scheme, host, port)
    if cache is not None:
        hit = cache.get(key)
        if hit is not None:
            return hit
    result = _discover_uncached(host, url, transport, timeout, resolve_inherits)
    if cache is not None:
        cache.put(key, result)
    return result


def _discover_uncached(host, url, transport, timeout, resolve_inherits) -> DiscoveryResult:
    try:
        final_url, doc = _fetch_json(url, transport, timeout)
        statute = _validate(doc, host, url)
        if statute.inherits and resolve_inherits:
            def fetch_parent(purl: str) -> Statute:
                phost = urlsplit(purl).hostname or ""
                try:
                    _, pdoc = _fetch_json(purl, transport, timeout)
                    return _validate(pdoc, phost, purl)
                except _Fail as e:
                    raise _Fail(FailureKind.InheritanceFailure,
                                f"parent {purl}: {e.kind.value}: {e.detail}", url)
            try:
                statute = resolve_inheritance(statute, url, fetch_parent)
            except InheritanceCycle as e:
                raise _Fail(FailureKind.InheritanceCycle, str(e), url)
        return Discovered(host, url, statute)
    except _Fail as e:
        return DiscoveryFailure(host, e.url, e.kind, e.detail)