utils/url_guard.py

"""
出站 URL 安全校验（防 SSRF）
============================

「模型设置」允许用户为 *custom*（OpenAI 兼容）提供商填写任意 ``base_url``，服务端随后
会带着 API 密钥向该地址发起请求（「测试」按钮与实际分析均会）。在公开多用户部署
（如 HF Spaces）上，任意用户可把 ``base_url`` 指向**内网 / 环回 / 云元数据**地址
（如 ``http://169.254.169.254/...``、``http://127.0.0.1``、``http://10.x``）借服务端
探测或读取内部资源——典型的 SSRF。

本模块对**用户提供**的 ``base_url`` 做校验：

- 必须为 ``http`` / ``https``（推荐 https）。
- 解析主机名后，拒绝任何解析到**环回 / 私有 / 链路本地 / 保留 / 多播**地址的目标
  （含字面量 IP 与经 DNS 解析的域名）。

注意：DNS 重绑定（TOCTOU）无法仅靠入口校验根治；本校验作为**纵深防御的第一道**，
显著抬高公开部署下的 SSRF 门槛。运营方在 ``CUSTOM_BASE_URL`` 环境变量中配置的内部
端点属可信来源，不经本校验（仅约束用户在 UI 中输入的值）。
"""

from __future__ import annotations

import ipaddress
import logging
import socket
from typing import Tuple
from urllib.parse import urlparse

logger = logging.getLogger(__name__)


def _ip_is_blocked(ip_str: str) -> bool:
    """判断 IP 是否落在禁止访问的网段（环回 / 私有 / 链路本地 / 保留 / 多播 / 未指定）。"""
    try:
        ip = ipaddress.ip_address(ip_str)
    except ValueError:
        return True  # 无法解析为 IP，保守拒绝
    return (
        ip.is_loopback
        or ip.is_private
        or ip.is_link_local
        or ip.is_reserved
        or ip.is_multicast
        or ip.is_unspecified
    )


def is_safe_public_url(url: str, *, require_https: bool = False) -> Tuple[bool, str]:
    """校验出站 URL 是否安全（非内网 / 非环回）。返回 ``(ok, reason)``。

    - 方案必须为 http/https（``require_https`` 为真时仅允许 https）。
    - 主机名解析出的**所有** IP 均不得落在被禁网段；任一命中即拒绝。
    - 解析失败、缺主机名等一律拒绝（``ok=False`` 并给出原因）。
    """
    if not url or not str(url).strip():
        return False, "URL 为空。"

    parsed = urlparse(str(url).strip())
    scheme = (parsed.scheme or "").lower()
    if scheme not in ("http", "https"):
        return False, f"不支持的协议：{scheme or '(空)'}（仅允许 http/https）。"
    if require_https and scheme != "https":
        return False, "出于安全要求，自定义端点必须使用 https。"

    host = parsed.hostname
    if not host:
        return False, "URL 缺少主机名。"

    # 字面量 IP：直接判定。
    try:
        ipaddress.ip_address(host)
        if _ip_is_blocked(host):
            return False, f"目标地址 {host} 属于内网 / 环回 / 保留网段，已拒绝。"
        return True, ""
    except ValueError:
        pass  # 不是字面量 IP，走 DNS 解析

    # 域名：解析全部 A/AAAA 记录，任一落在禁止网段即拒绝。
    try:
        infos = socket.getaddrinfo(host, parsed.port or (443 if scheme == "https" else 80),
                                   proto=socket.IPPROTO_TCP)
    except (socket.gaierror, OSError) as exc:
        return False, f"无法解析主机 {host}：{exc}"

    resolved = {info[4][0] for info in infos}
    if not resolved:
        return False, f"主机 {host} 未解析到任何地址。"
    for ip_str in resolved:
        if _ip_is_blocked(ip_str):
            return False, f"主机 {host} 解析到内网 / 环回地址 {ip_str}，已拒绝。"
    return True, ""


__all__ = ["is_safe_public_url"]