xiaoxia-saas/packages/domain/url_security.py

"""URL 安全校验纯逻辑 — SSRF 防护.

纯函数模块，无网络/文件 IO，无环境变量依赖，所有配置通过参数传入。
供 packages/shared/url_security.py 作为薄包装调用，也可直接用于单测。

防护要点（纯逻辑部分）：
1. Scheme 白名单校验
2. 内部主机名拦截（字符串匹配）
3. 端口白名单校验
4. IP 格式 SSRF 检查（回环/私有/链路本地/组播/未指定/保留）
5. 可信域名匹配（支持子域名）
6. 文件头魔数校验（接收 bytes）
"""

from __future__ import annotations

import ipaddress
from urllib.parse import urlparse


class UrlSecurityError(ValueError):
    """URL 安全校验失败."""

    pass


# ── 常量 ────────────────────────────────────────────────────────────────────

ALLOWED_SCHEMES = frozenset({"http", "https"})
ALLOWED_PORTS = frozenset({80, 443})
MAX_URL_LENGTH = 2048

# 已知内部/敏感主机名集合
INTERNAL_HOSTNAMES = frozenset(
    {
        "localhost",
        "localhost.localdomain",
        "ip6-localhost",
        "ip6-loopback",
        "metadata",
        "metadata.google.internal",
        "169.254.169.254",
    }
)

# 内网域名后缀
INTERNAL_DOMAIN_SUFFIXES = (".local", ".internal", ".localdomain")

# 文件魔数表 — key: MIME, value: list of 签名组，每组内所有 (offset, bytes) 都匹配才算命中
MAGIC_NUMBERS: dict[str, list[list[tuple[int, bytes]]]] = {
    # 音频
    "audio/mpeg": [
        [(0, b"ID3")],
        [(0, b"\xff\xfb")],
        [(0, b"\xff\xf3")],
        [(0, b"\xff\xf2")],
        [(0, b"\xff\xfa")],
        [(0, b"\xff\xf9")],
    ],
    "audio/wav": [[(0, b"RIFF"), (8, b"WAVE")]],
    "audio/x-wav": [[(0, b"RIFF"), (8, b"WAVE")]],
    "audio/ogg": [[(0, b"OggS")]],
    "application/ogg": [[(0, b"OggS")]],
    "audio/flac": [[(0, b"fLaC")]],
    "audio/aac": [[(0, b"\xff\xf1")], [(0, b"\xff\xf9")]],
    "audio/aacp": [[(0, b"\xff\xf1")], [(0, b"\xff\xf9")]],
    "audio/mp4": [[(4, b"ftyp")]],
    "audio/x-m4a": [[(4, b"ftyp")]],
    # 视频
    "video/mp4": [[(4, b"ftyp")]],
    "video/quicktime": [[(4, b"ftyp")]],
    "video/x-matroska": [[(0, b"\x1a\x45\xdf\xa3")]],
    "video/webm": [[(0, b"\x1a\x45\xdf\xa3")]],
    "video/x-msvideo": [[(0, b"RIFF"), (8, b"AVI ")]],
    # 图片
    "image/jpeg": [[(0, b"\xff\xd8\xff")]],
    "image/png": [[(0, b"\x89PNG\r\n\x1a\n")]],
    "image/gif": [[(0, b"GIF87a")], [(0, b"GIF89a")]],
    "image/webp": [[(0, b"RIFF"), (8, b"WEBP")]],
    "image/bmp": [[(0, b"BM")]],
}

ALLOWED_AUDIO_MIME_TYPES = frozenset(
    {
        "audio/mpeg",
        "audio/mp3",
        "audio/wav",
        "audio/x-wav",
        "audio/pcm",
        "audio/ogg",
        "audio/opus",
        "audio/flac",
        "audio/aac",
        "audio/m4a",
        "audio/x-m4a",
        "audio/mp4",
        "application/octet-stream",
    }
)

ALLOWED_VIDEO_MIME_TYPES = frozenset(
    {
        "video/mp4",
        "video/quicktime",
        "video/x-matroska",
        "video/webm",
        "video/avi",
        "video/x-msvideo",
        "video/mpeg",
        "application/octet-stream",
    }
)

ALLOWED_IMAGE_MIME_TYPES = frozenset(
    {
        "image/jpeg",
        "image/png",
        "image/gif",
        "image/webp",
        "image/bmp",
    }
)


# ── 主机名 / 域名校验 ──────────────────────────────────────────────────────


def check_internal_hostname(hostname: str) -> None:
    """检查主机名是否为内部/敏感主机名，是则抛出 UrlSecurityError.

    检查内容：
    - 精确匹配 INTERNAL_HOSTNAMES 集合
    - 后缀匹配 INTERNAL_DOMAIN_SUFFIXES（.local/.internal/.localdomain）
    """
    hostname_lower = hostname.lower()
    if hostname_lower in INTERNAL_HOSTNAMES:
        raise UrlSecurityError(f"禁止访问内部主机名: {hostname}")
    if hostname_lower.endswith(INTERNAL_DOMAIN_SUFFIXES):
        raise UrlSecurityError(f"禁止访问内网域名: {hostname}")


def is_trusted_domain(hostname: str, trusted_domains: set[str]) -> bool:
    """检查域名是否在可信白名单中（支持子域名匹配）.

    匹配规则：
    - 精确匹配
    - 子域名匹配（hostname 以 .domain 结尾）

    Args:
        hostname: 待检查的主机名
        trusted_domains: 可信域名集合，为空表示不限制
    """
    if not trusted_domains:
        return True
    hostname_lower = hostname.lower()
    if hostname_lower in {d.lower() for d in trusted_domains}:
        return True
    for domain in trusted_domains:
        if hostname_lower.endswith("." + domain.lower()):
            return True
    return False


# ── IP SSRF 检查 ───────────────────────────────────────────────────────────


def check_ssrf_ip(ip_str: str) -> None:
    """检查 IP 地址是否存在 SSRF 风险，有风险则抛出 UrlSecurityError.

    检查项：回环、私有、链路本地、组播、未指定、保留地址。

    Args:
        ip_str: IP 地址字符串（IPv4 或 IPv6）

    Raises:
        UrlSecurityError: IP 属于 SSRF 风险范围
        ValueError: ip_str 不是合法 IP 地址（调用方应自行捕获处理）
    """
    ip_obj = ipaddress.ip_address(ip_str)
    if ip_obj.is_loopback:
        raise UrlSecurityError(f"禁止访问回环地址: {ip_obj}")
    if ip_obj.is_link_local:
        raise UrlSecurityError(f"禁止访问链路本地地址: {ip_obj}")
    if ip_obj.is_unspecified:
        raise UrlSecurityError(f"禁止访问未指定地址: {ip_obj}")
    if ip_obj.is_multicast:
        raise UrlSecurityError(f"禁止访问组播地址: {ip_obj}")
    if ip_obj.is_reserved:
        raise UrlSecurityError(f"禁止访问保留地址: {ip_obj}")
    if ip_obj.is_private:
        raise UrlSecurityError(f"禁止访问内网地址: {ip_obj}")


def is_ip_address(hostname: str) -> bool:
    """判断主机名是否为 IP 地址格式（IPv4 或 IPv6）."""
    try:
        ipaddress.ip_address(hostname)
        return True
    except ValueError:
        return False


# ── URL 基础校验 ───────────────────────────────────────────────────────────


def validate_url_basic(
    url: str,
    *,
    trusted_domains: set[str] | None = None,
    allow_direct_ip: bool = False,
) -> str:
    """URL 基础安全校验（纯逻辑，不含 DNS 解析）.

    校验项：
    1. URL 非空 & 长度限制
    2. Scheme 白名单
    3. 主机名存在性
    4. 内部主机名拦截
    5. 端口白名单
    6. 直接 IP 访问限制
    7. IP 格式 SSRF 检查（如果 hostname 是 IP）
    8. 可信域名白名单（如果配置了）

    注意：域名格式的 SSRF 检查需要 DNS 解析，不在本函数范围内。

    Args:
        url: 待校验 URL
        trusted_domains: 可信域名白名单，None/空表示不限制
        allow_direct_ip: 是否允许直接 IP 访问

    Returns:
        原始 URL（校验通过）

    Raises:
        UrlSecurityError: 校验失败
    """
    if not url:
        raise UrlSecurityError("URL 为空")
    if len(url) > MAX_URL_LENGTH:
        raise UrlSecurityError(f"URL 过长 ({len(url)} > {MAX_URL_LENGTH})")

    try:
        parsed = urlparse(url)
    except Exception as e:
        raise UrlSecurityError(f"URL 解析失败: {e}") from e

    # Scheme
    if not parsed.scheme or parsed.scheme.lower() not in ALLOWED_SCHEMES:
        raise UrlSecurityError(f"不允许的 URL scheme: {parsed.scheme}")

    # Hostname
    hostname = parsed.hostname
    if not hostname:
        raise UrlSecurityError("URL 缺少主机名")

    # 内部主机名前置拦截
    check_internal_hostname(hostname)

    # 端口
    port = parsed.port
    if port is not None and port not in ALLOWED_PORTS:
        raise UrlSecurityError(f"不允许的端口: {port}")

    # IP 格式检查 & SSRF
    if is_ip_address(hostname):
        if not allow_direct_ip:
            raise UrlSecurityError(f"禁止直接 IP 访问: {hostname}")
        check_ssrf_ip(hostname)

    # 可信域名白名单
    if trusted_domains and not is_trusted_domain(hostname, trusted_domains):
        raise UrlSecurityError(f"域名不在可信白名单中: {hostname}")

    return url


def is_url_basic_safe(
    url: str,
    *,
    trusted_domains: set[str] | None = None,
    allow_direct_ip: bool = False,
) -> bool:
    """便捷函数：基础安全检查，不抛异常，返回 bool."""
    try:
        validate_url_basic(url, trusted_domains=trusted_domains, allow_direct_ip=allow_direct_ip)
        return True
    except UrlSecurityError:
        return False


# ── 魔数校验 ───────────────────────────────────────────────────────────────


def validate_magic_number(header_bytes: bytes, allowed_mime_types: set[str]) -> None:
    """校验文件头魔数是否与允许的 MIME 类型匹配（纯函数）.

    读取 header_bytes，与 allowed_mime_types 对应格式的魔数逐一比对，
    任一类型匹配即通过；全部不匹配则抛出 UrlSecurityError。

    Args:
        header_bytes: 文件头字节（建议至少 256 字节）
        allowed_mime_types: 允许的 MIME 类型集合

    Raises:
        UrlSecurityError: 文件魔数与所有允许类型均不匹配
    """
    # 收集所有允许类型对应的魔数签名
    signatures: list[list[tuple[int, bytes]]] = []
    for mime in allowed_mime_types:
        sigs = MAGIC_NUMBERS.get(mime)
        if sigs:
            signatures.extend(sigs)

    # 没有已知魔数的 MIME，跳过不阻断
    if not signatures:
        return

    if not header_bytes:
        raise UrlSecurityError("文件为空，无法校验格式")

    # 任一签名匹配即通过
    for sig in signatures:
        match = True
        for offset, expected in sig:
            if offset + len(expected) > len(header_bytes):
                match = False
                break
            if header_bytes[offset : offset + len(expected)] != expected:
                match = False
                break
        if match:
            return

    raise UrlSecurityError(
        f"文件魔数与允许的 MIME 类型不匹配，"
        f"允许类型: {sorted(allowed_mime_types)}，"
        f"文件头前16字节: {header_bytes[:16].hex()}"
    )