"""SSRF-guarded outbound HTTP for autodiscovery (SPEC §10.6).

All autodiscovery HTTP is HTTPS-only, 5 s timeout, response capped at 256 KB, no
off-domain redirects, and the target host must resolve exclusively to public IPs
(reject private/loopback/link-local/reserved ranges). A residual TOCTOU gap
(DNS rebinding between our resolve and httpx's own resolve) is accepted for M1;
a pinned-IP transport can close it later.
"""

from __future__ import annotations

import asyncio
import ipaddress
import socket
from urllib.parse import urlsplit

import httpx

HTTP_TIMEOUT = 5.0
MAX_RESPONSE_BYTES = 256 * 1024


class SsrfError(Exception):
    """Raised when a URL fails the SSRF policy (bad scheme, private IP, oversize)."""


def _is_public_ip(ip: str) -> bool:
    addr = ipaddress.ip_address(ip)
    return not (
        addr.is_private
        or addr.is_loopback
        or addr.is_link_local
        or addr.is_multicast
        or addr.is_reserved
        or addr.is_unspecified
    )


async def _resolve_all(host: str) -> list[str]:
    loop = asyncio.get_running_loop()
    infos = await loop.getaddrinfo(host, None, proto=socket.IPPROTO_TCP)
    return [info[4][0] for info in infos]


async def assert_host_is_public(host: str) -> None:
    """Resolve ``host`` and reject if any address is non-public (SPEC §10.6)."""
    try:
        addresses = await _resolve_all(host)
    except socket.gaierror as exc:
        raise SsrfError(f"cannot resolve {host}") from exc
    if not addresses:
        raise SsrfError(f"no addresses for {host}")
    for ip in addresses:
        if not _is_public_ip(ip):
            raise SsrfError(f"{host} resolves to non-public address {ip}")


async def guarded_get(url: str) -> str | None:
    """HTTPS GET with the full SSRF policy applied.

    Returns the body text on 200, or ``None`` when the resource is absent
    (any non-200 response). Raises :class:`SsrfError` on a policy violation
    (bad scheme, private IP, oversize); lets httpx network errors propagate.
    """
    parts = urlsplit(url)
    if parts.scheme != "https":
        raise SsrfError("only https is permitted")
    if not parts.hostname:
        raise SsrfError("missing host")
    await assert_host_is_public(parts.hostname)

    async with (
        httpx.AsyncClient(
            follow_redirects=False,  # off-domain redirects are forbidden (§10.6)
            timeout=HTTP_TIMEOUT,
            headers={"User-Agent": "webmail-autodiscovery/1.0"},
        ) as client,
        client.stream("GET", url) as resp,
    ):
        if resp.status_code != 200:
            return None
        chunks: list[bytes] = []
        total = 0
        async for chunk in resp.aiter_bytes():
            total += len(chunk)
            if total > MAX_RESPONSE_BYTES:
                raise SsrfError("response exceeds 256 KB cap")
            chunks.append(chunk)
    return b"".join(chunks).decode("utf-8", "replace")