ai-kefu-admin/app/platform/site_import/safe_fetch.py

"""site_import 安全抓取原语（spec §3.1）。

连接层 SSRF 校验：自定义 httpcore 后端在 connect_tcp 处做
anyio.getaddrinfo 解析（异步——单进程零同步 IO 纪律）→ IP 禁段校验 →
只连已验 IP 字面量。DNS rebinding 的 TOCTOU 由「解析与连接使用同一
后端的同一结果」闭合。禁止改写 URL host 为 IP（会破坏 TLS 证书校验）——
SNI/证书由 httpcore 以原 hostname 包 TLS 保持正确
（httpcore connection.py: start_tls(server_hostname=origin.host)）。

私有 API 依赖（httpx 0.28 / httpcore 1.0.9 钉位，tests/test_safe_fetch.py
版本探针守卫，升级破坏时立即红）：
- httpx.AsyncHTTPTransport._pool 属性替换注入校验连接池
- httpcore.AsyncConnectionPool(network_backend=...)
- httpcore._backends.anyio.AnyIOStream（anyio 流 → httpcore 流包装）
- httpcore._exceptions.ConnectError / ConnectTimeout / map_exceptions
"""
import ipaddress
import socket
import ssl
import urllib.parse

import anyio
import certifi
import httpcore
import httpx
from httpcore._backends.anyio import AnyIOStream
from httpcore._exceptions import ConnectError, ConnectTimeout, map_exceptions

BLOCKED_NETS = [ipaddress.ip_network(n) for n in [
    "127.0.0.0/8", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16",
    "169.254.0.0/16", "0.0.0.0/8", "100.64.0.0/10",
    "224.0.0.0/4", "240.0.0.0/4", "198.18.0.0/15",
    "::1/128", "fc00::/7", "fe80::/10", "::ffff:0:0/96",
    "::/128", "64:ff9b::/96",
]]
ALLOWED_PORTS = {80, 443}
MAX_REDIRECTS = 5
USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"
)


class BlockedAddressError(Exception):
    """目标解析到禁段（或端口非法）——连接前拒绝。"""


class FetchError(Exception):
    """抓取失败（URL 非法 / 端口禁 / 重定向超限 / 请求层错误）。"""


def is_blocked_ip(ip: str) -> bool:
    """禁段判定；非法 IP 串一律视为禁（fail-closed）。"""
    try:
        addr = ipaddress.ip_address(ip)
    except ValueError:
        return True
    return any(addr in net for net in BLOCKED_NETS)


def _is_blocked_literal(hostname: str) -> bool:
    """URL 层的 host 禁段预检——只拒「IP 字面量落在禁段」的 host。

    普通域名（非 IP 字面量）此处放行，交由连接层解析后按 is_blocked_ip
    校验（is_blocked_ip 对非 IP 串 fail-closed 为 True，那是给解析结果
    用的语义，不能套在域名上）。
    """
    try:
        addr = ipaddress.ip_address(hostname)
    except ValueError:
        return False
    return is_blocked_ip(str(addr))


def _sockaddr_ip(raw) -> str:
    """getaddrinfo 的 sockaddr[0] 兼容 str / 打包 bytes 两种形态（IPv6 bytes）。"""
    if isinstance(raw, (bytes, bytearray)):
        return str(ipaddress.ip_address(bytes(raw)))
    return str(raw)


class ValidatingBackend(httpcore.AsyncNetworkBackend):
    """connect_tcp 钩子：anyio.getaddrinfo 解析 → 逐候选禁段校验 → 只连已验 IP 字面量。

    评审折入的两个实现陷阱：
    1. anyio 不重导出 socket 常量——AF_UNSPEC/SOCK_STREAM 用 socket 模块；
    2. 默认后端 anyio.connect_tcp(remote_host=host) 会对 hostname 内部重新做
       DNS 解析——把 rebinding TOCTOU 重新打开。本实现只把「已验 IP 字面量」
       递给连接调用（anyio 对 IP 字面量跳过 DNS 直连），原 hostname 不再进
       任何会做 DNS 的路径。connect_unix_socket 不实现（基类默认拒绝）。
    """

    async def connect_tcp(
        self,
        host: str,
        port: int,
        timeout: float | None = None,
        local_address: str | None = None,
        socket_options=None,
    ):
        if port not in ALLOWED_PORTS:
            raise BlockedAddressError(
                f"port {port} not in allowed {sorted(ALLOWED_PORTS)}")
        if socket_options is None:
            socket_options = []
        exc_map = {
            TimeoutError: ConnectTimeout,
            OSError: ConnectError,
            anyio.BrokenResourceError: ConnectError,
        }
        with map_exceptions(exc_map):
            infos = await anyio.getaddrinfo(
                host, port, family=socket.AF_UNSPEC, type=socket.SOCK_STREAM)
            validated: str | None = None
            for _family, _type, _proto, _canon, sockaddr in infos:
                ip = _sockaddr_ip(sockaddr[0])
                if is_blocked_ip(ip):
                    raise BlockedAddressError(
                        f"{host!r} resolves to blocked address {ip}")
                if validated is None:
                    validated = ip
            if validated is None:  # pragma: no cover - getaddrinfo 空结果防御
                raise BlockedAddressError(f"{host!r} resolved to no addresses")
            with anyio.fail_after(timeout):
                stream = await anyio.connect_tcp(
                    remote_host=validated, remote_port=port, local_host=local_address)
                for option in socket_options:
                    stream._raw_socket.setsockopt(*option)  # type: ignore[attr-defined]
        return AnyIOStream(stream)

    async def sleep(self, seconds: float) -> None:
        await anyio.sleep(seconds)


def _make_client(timeout: float = 10.0) -> httpx.AsyncClient:
    """校验后端注入的 client（私有口 _pool 替换——版本探针测试钉住）。"""
    transport = httpx.AsyncHTTPTransport()
    transport._pool = httpcore.AsyncConnectionPool(
        ssl_context=ssl.create_default_context(cafile=certifi.where()),
        network_backend=ValidatingBackend(),
    )
    return httpx.AsyncClient(
        transport=transport,
        timeout=timeout,
        follow_redirects=False,
        headers={"User-Agent": USER_AGENT},
    )


def _validate_url(url: str) -> None:
    """URL 层预检：scheme 白名单 + host 禁段 + 端口白名单（连接层校验的前置快速失败）。"""
    try:
        parsed = urllib.parse.urlsplit(url)
        port = parsed.port
    except ValueError as exc:
        raise FetchError(f"invalid URL: {url!r}") from exc
    if parsed.scheme not in ("http", "https"):
        raise FetchError(f"scheme not allowed: {parsed.scheme!r} ({url!r})")
    if not parsed.hostname:
        raise FetchError(f"URL without host: {url!r}")
    if _is_blocked_literal(parsed.hostname):
        raise FetchError(f"URL host in blocked range: {parsed.hostname}")
    if port is None:
        port = 443 if parsed.scheme == "https" else 80
    if port not in ALLOWED_PORTS:
        raise FetchError(f"port {port} not allowed: {url!r}")


async def fetch(url: str, timeout: float = 10.0) -> httpx.Response:
    """单页安全抓取：URL 预检 + 逐跳手动重定向（≤5 跳，每跳重过完整校验链）。

    双层防线：本函数的 URL 层预检挡住字面量私网/非标端口；ValidatingBackend
    在连接层对每一跳的 DNS 解析结果做禁段校验（防 rebinding）。网络层错误
    统一包装为 FetchError（保留 __cause__）。
    """
    _validate_url(url)
    client = _make_client(timeout)
    current = url
    redirects = 0
    try:
        while True:
            response = await client.get(current)
            if not response.has_redirect_location:
                return response
            redirects += 1
            if redirects > MAX_REDIRECTS:
                raise FetchError(
                    f"too many redirects (>{MAX_REDIRECTS}) at {current!r}")
            current = urllib.parse.urljoin(current, response.headers["location"])
            _validate_url(current)
    except (httpx.HTTPError, httpx.InvalidURL, BlockedAddressError, OSError) as exc:
        raise FetchError(f"fetch failed for {url!r}: {exc}") from exc
    finally:
        await client.aclose()