xiaoxia-saas/packages/shared/url_security.py

"""URL 安全校验工具 — SSRF 防护.

统一的外部 URL 安全校验方案，覆盖所有渲染管线和 TTS 中的外部下载场景。
放在 packages/shared/ 作为单一来源，worker 和 application 层都可引用。

防护要点：
1. Scheme 白名单：仅允许 http/https
2. 主机 SSRF 防护：禁止内网 IP、回环地址、链路本地地址、元数据服务
3. 端口白名单：仅允许 80/443（标准 HTTP/HTTPS）
4. 域名校验：禁止 IP 直接访问（除非在白名单中）
5. 重定向防护：手动跟随重定向，每次跳转前重新校验目标 URL
6. 文件大小限制：流式下载，超过上限立即中断
7. MIME 类型白名单：可选的内容类型校验
"""

from __future__ import annotations

import ipaddress
import logging
import os
import socket
import urllib.error
import urllib.request
from urllib.parse import urljoin, urlparse

logger = logging.getLogger(__name__)

# 允许的 URL scheme
ALLOWED_SCHEMES = {"http", "https"}

# 允许的端口（标准 HTTP/HTTPS）
ALLOWED_PORTS = {80, 443}

# 可信域名白名单（可根据实际 OSS/CDN 域名配置）
# 从环境变量读取，格式："oss-cn-hangzhou.aliyuncs.com,cdn.example.com"
# 默认空表示所有公网域名都允许，但仍会做 SSRF 检查
TRUSTED_DOMAINS: set[str] = set()
_env_trusted = os.environ.get("URL_SECURITY_TRUSTED_DOMAINS", "")
if _env_trusted:
    TRUSTED_DOMAINS = {d.strip() for d in _env_trusted.split(",") if d.strip()}

# 是否允许 IP 直接访问（默认禁止，防止绕过 DNS 校验）
ALLOW_DIRECT_IP = os.environ.get("URL_SECURITY_ALLOW_DIRECT_IP", "false").lower() == "true"

# 最大 URL 长度
MAX_URL_LENGTH = 2048

# 单次下载最大文件大小（默认 200MB）
DEFAULT_MAX_DOWNLOAD_SIZE = int(os.environ.get("URL_SECURITY_MAX_DOWNLOAD_MB", "200")) * 1024 * 1024

# 允许的音频 MIME 类型白名单
ALLOWED_AUDIO_MIME_TYPES = {
    "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",  # 兼容一些 CDN 返回通用类型
}

# 允许的视频 MIME 类型白名单
ALLOWED_VIDEO_MIME_TYPES = {
    "video/mp4",
    "video/quicktime",
    "video/x-matroska",
    "video/webm",
    "video/avi",
    "video/x-msvideo",
    "video/mpeg",
    "application/octet-stream",
}

# 允许的图片 MIME 类型白名单
ALLOWED_IMAGE_MIME_TYPES = {
    "image/jpeg",
    "image/png",
    "image/gif",
    "image/webp",
    "image/bmp",
}

# 下载块大小
_DOWNLOAD_CHUNK_SIZE = 8192

# 最大重定向次数
_MAX_REDIRECTS = 5

# 文件魔数（文件头签名）表 — 用于 MIME 白名单校验后的二次真实性校验
# key: MIME 类型，value: 签名列表，任一签名匹配即通过
# 每条签名: list of (offset, bytes)，所有条目都匹配才算该签名命中（支持多处联合匹配如 RIFF+WAVE）
_MAGIC_NUMBERS: dict[str, list[list[tuple[int, bytes]]]] = {
    # ── 音频 ──
    "audio/mpeg": [
        [(0, b"ID3")],  # ID3v2 标签
        [(0, b"\xff\xfb")],  # MPEG1 Layer3
        [(0, b"\xff\xf3")],  # MPEG2 Layer3
        [(0, b"\xff\xf2")],  # MPEG2.5 Layer3
        [(0, b"\xff\xfa")],  # MPEG1 Layer2
        [(0, b"\xff\xf9")],  # 其他 MPEG ADTS
    ],
    "audio/wav": [
        [(0, b"RIFF"), (8, b"WAVE")],  # RIFF + 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")],  # ADTS MPEG-4
        [(0, b"\xff\xf9")],  # ADTS MPEG-2
    ],
    "audio/aacp": [
        [(0, b"\xff\xf1")],
        [(0, b"\xff\xf9")],
    ],
    "audio/mp4": [
        [(4, b"ftyp")],  # ISO Base Media (M4A)
    ],
    "audio/x-m4a": [
        [(4, b"ftyp")],
    ],
    # ── 视频 ──
    "video/mp4": [
        [(4, b"ftyp")],  # ISO Base Media (MP4)
    ],
    "video/quicktime": [
        [(4, b"ftyp")],
    ],
    "video/x-matroska": [
        [(0, b"\x1a\x45\xdf\xa3")],  # EBML header
    ],
    "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")],
    ],
}

# 魔数校验最大读取字节数（文件头）
_MAGIC_CHECK_READ_SIZE = 256


def _validate_magic_number(file_path: str, allowed_mime_types: set[str]) -> None:
    """校验文件头魔数是否与允许的 MIME 类型匹配.

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

    仅当 allowed_mime_types 非空时执行；空文件视为不匹配。

    Args:
        file_path: 本地文件路径
        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

    try:
        with open(file_path, "rb") as f:
            header = f.read(_MAGIC_CHECK_READ_SIZE)
    except OSError as e:
        raise UrlSecurityError(f"读取文件头失败: {e}") from e

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

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

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


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

    pass


class NoRedirectHandler(urllib.request.HTTPRedirectHandler):
    """禁止自动重定向的 handler，用于手动控制重定向以做安全校验."""

    def redirect_request(self, req, fp, code, msg, headers, newurl):  # noqa: N802
        return None


def validate_url_safety(url: str, *, purpose: str = "download") -> str:
    """校验 URL 安全性，返回标准化后的 URL（供下游使用）.

    Args:
        url: 待校验的 URL
        purpose: 用途描述（用于日志），如 "bgm_download"、"tts_download"

    Returns:
        标准化后的 URL

    Raises:
        UrlSecurityError: URL 不安全
    """
    if not url:
        raise UrlSecurityError("URL 为空")

    if len(url) > MAX_URL_LENGTH:
        raise UrlSecurityError(f"URL 过长 ({len(url)} > {MAX_URL_LENGTH})")

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

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

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

    # 2.1 常见内网主机名前置拦截（防止 DNS rebinding 绕过）
    _check_internal_hostnames(hostname)

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

    # 4. SSRF 防护 - 解析 IP 并检查
    try:
        # 先判断是否是 IP 地址
        ip_obj = None
        try:
            ip_obj = ipaddress.ip_address(hostname)
        except ValueError:
            pass  # 不是 IP，继续走域名解析

        if ip_obj is not None:
            # 是直接 IP 访问
            if not ALLOW_DIRECT_IP and not _is_trusted_ip(ip_obj):
                raise UrlSecurityError(f"禁止直接 IP 访问: {hostname}")
            _check_ssrf_ip(ip_obj)
        else:
            # 域名 — 解析 DNS 检查 SSRF
            _check_ssrf_domain(hostname)
    except UrlSecurityError:
        raise
    except Exception as e:
        logger.warning("URL 安全校验异常: url=%s purpose=%s error=%s", url[:80], purpose, e)
        raise UrlSecurityError(f"URL 安全校验异常: {e}") from e

    # 5. 可信域名校验（如果配置了白名单）
    if TRUSTED_DOMAINS and not _is_trusted_domain(hostname):
        raise UrlSecurityError(f"域名不在可信白名单中: {hostname}")

    logger.debug("URL 安全校验通过: url=%s purpose=%s", url[:80], purpose)
    return url


def _check_internal_hostnames(hostname: str) -> None:
    """前置检查常见内网/敏感主机名，防止 DNS 解析层绕过."""
    hostname_lower = hostname.lower()
    internal_hostnames = {
        "localhost",
        "localhost.localdomain",
        "ip6-localhost",
        "ip6-loopback",
        "metadata",
        "metadata.google.internal",
        "169.254.169.254",  # 云元数据服务
    }
    if hostname_lower in internal_hostnames:
        raise UrlSecurityError(f"禁止访问内部主机名: {hostname}")

    # 检查以 .local / .internal 结尾的主机名
    if hostname_lower.endswith((".local", ".internal", ".localdomain")):
        raise UrlSecurityError(f"禁止访问内网域名: {hostname}")


def _check_ssrf_ip(ip_obj: ipaddress.IPv4Address | ipaddress.IPv6Address) -> None:
    """检查 IP 是否属于 SSRF 风险范围."""
    # 回环地址
    if ip_obj.is_loopback:
        raise UrlSecurityError(f"禁止访问回环地址: {ip_obj}")

    # 私有地址（内网）
    if ip_obj.is_private:
        raise UrlSecurityError(f"禁止访问内网地址: {ip_obj}")

    # 链路本地地址
    if ip_obj.is_link_local:
        raise UrlSecurityError(f"禁止访问链路本地地址: {ip_obj}")

    # 组播地址
    if ip_obj.is_multicast:
        raise UrlSecurityError(f"禁止访问组播地址: {ip_obj}")

    # 未指定地址（0.0.0.0 / ::）
    if ip_obj.is_unspecified:
        raise UrlSecurityError(f"禁止访问未指定地址: {ip_obj}")

    # 保留地址
    if ip_obj.is_reserved:
        raise UrlSecurityError(f"禁止访问保留地址: {ip_obj}")


def _check_ssrf_domain(hostname: str) -> None:
    """对域名做 DNS 解析并检查所有解析结果的 IP 是否安全.

    注意：这不能完全防止 DNS rebinding，但能防御大部分 SSRF 场景。
    """
    try:
        # 解析所有地址
        infos = socket.getaddrinfo(hostname, None)
        if not infos:
            raise UrlSecurityError(f"域名解析失败: {hostname}")

        for info in infos:
            ip_str = info[4][0]
            try:
                ip_obj = ipaddress.ip_address(ip_str)
                _check_ssrf_ip(ip_obj)
            except ValueError:
                # 无法解析为 IP，跳过（不应该发生）
                continue
    except socket.gaierror as e:
        raise UrlSecurityError(f"域名解析失败: {hostname} ({e})") from e


def _is_trusted_ip(ip_obj: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
    """检查 IP 是否在可信列表中（目前通过环境变量配置域名，IP 级信任暂不开放）."""
    return False


def _is_trusted_domain(hostname: str) -> bool:
    """检查域名是否在可信白名单中（支持子域名匹配）."""
    hostname_lower = hostname.lower()
    if hostname_lower in TRUSTED_DOMAINS:
        return True
    # 检查子域名
    for domain in TRUSTED_DOMAINS:
        if hostname_lower.endswith("." + domain.lower()):
            return True
    return False


def is_url_safe(url: str, *, purpose: str = "download") -> bool:
    """便捷函数：检查 URL 是否安全，不抛异常."""
    try:
        validate_url_safety(url, purpose=purpose)
        return True
    except UrlSecurityError:
        return False


# ── 安全下载 ───────────────────────────────────────────────────────────────────


def safe_download_file(
    url: str,
    dest_path: str,
    *,
    purpose: str = "download",
    max_size: int = DEFAULT_MAX_DOWNLOAD_SIZE,
    allowed_mime_types: set[str] | None = None,
    timeout: float = 60.0,
) -> int:
    """安全下载 URL 到本地文件。

    包含防护：
    - SSRF 校验（初始 URL + 每次重定向后都校验）
    - 重定向次数限制 + 手动跟随（避免重定向绕过 SSRF）
    - 文件大小限制（流式读取，超过立即中断）
    - MIME 类型白名单（可选）
    - 文件头魔数校验（配合 MIME 白名单做二次真实性校验）

    Args:
        url: 下载 URL
        dest_path: 目标文件路径
        purpose: 用途描述（日志用）
        max_size: 最大下载字节数，超过则中断并抛出 UrlSecurityError
        allowed_mime_types: 允许的 Content-Type 集合，None 表示不校验
        timeout: 单次请求超时（秒）

    Returns:
        实际下载的字节数

    Raises:
        UrlSecurityError: 安全校验失败
    """
    current_url = url
    redirect_count = 0
    total_bytes = 0

    # 使用不自动跟随重定向的 opener
    no_redirect_opener = urllib.request.build_opener(NoRedirectHandler())

    while True:
        # 每次请求前都做 SSRF 校验（重定向目标也会校验）
        validate_url_safety(current_url, purpose=purpose)

        req = urllib.request.Request(current_url, method="GET")
        req.add_header("User-Agent", "xiaoxia-saas-worker/1.0")

        try:
            resp = no_redirect_opener.open(req, timeout=timeout)  # nosec B310
        except urllib.error.HTTPError as e:
            # 3xx 重定向
            if 300 <= e.code < 400 and e.headers.get("Location"):
                if redirect_count >= _MAX_REDIRECTS:
                    raise UrlSecurityError(f"重定向次数超过限制 ({_MAX_REDIRECTS})") from e
                redirect_count += 1
                current_url = urljoin(current_url, e.headers["Location"])
                continue
            raise UrlSecurityError(f"HTTP 错误: {e.code} {e.reason}") from e
        except urllib.error.URLError as e:
            raise UrlSecurityError(f"URL 错误: {e.reason}") from e

        try:
            # Content-Type 校验
            if allowed_mime_types is not None:
                content_type = resp.headers.get("Content-Type", "").split(";")[0].strip().lower()
                if content_type and content_type not in allowed_mime_types:
                    raise UrlSecurityError(
                        f"不允许的 Content-Type: {content_type}, " f"允许: {sorted(allowed_mime_types)}"
                    )

            # Content-Length 预检
            content_length = resp.headers.get("Content-Length")
            if content_length and int(content_length) > max_size:
                raise UrlSecurityError(f"文件过大: {content_length} bytes > {max_size} bytes 上限")

            # 流式下载，实时检查大小
            with open(dest_path, "wb") as f:
                while True:
                    chunk = resp.read(_DOWNLOAD_CHUNK_SIZE)
                    if not chunk:
                        break
                    total_bytes += len(chunk)
                    if total_bytes > max_size:
                        raise UrlSecurityError(f"下载超过大小限制: {total_bytes} bytes > {max_size} bytes")
                    f.write(chunk)

            # 文件头魔数校验（MIME 白名单基础上的二次真实性校验）
            if allowed_mime_types is not None:
                _validate_magic_number(dest_path, allowed_mime_types)

            return total_bytes
        finally:
            resp.close()


def safe_download_bytes(
    url: str,
    *,
    purpose: str = "download",
    max_size: int = DEFAULT_MAX_DOWNLOAD_SIZE,
    allowed_mime_types: set[str] | None = None,
    timeout: float = 60.0,
) -> bytes:
    """安全下载 URL 并返回字节内容。

    防护同 safe_download_file，但结果返回在内存中（适合小文件）。
    """
    import tempfile

    fd, tmp_path = tempfile.mkstemp()
    os.close(fd)

    try:
        safe_download_file(
            url,
            tmp_path,
            purpose=purpose,
            max_size=max_size,
            allowed_mime_types=allowed_mime_types,
            timeout=timeout,
        )
        with open(tmp_path, "rb") as f:
            return f.read()
    finally:
        try:
            os.unlink(tmp_path)
        except OSError:
            pass