xiaoxia-saas/packages/shared/url_security.py

"""URL 安全校验工具 — SSRF 防护（薄包装层）.

本文件保留原有对外 API，纯逻辑部分委托给 packages/domain/url_security.py。
新增了 DNS 解析、文件下载、环境变量配置等有副作用的逻辑。

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

from __future__ import annotations

import logging
import os
import socket
import urllib.error
import urllib.request
from urllib.parse import urljoin

from packages.domain.url_security import (
    ALLOWED_AUDIO_MIME_TYPES,
    ALLOWED_IMAGE_MIME_TYPES,
    ALLOWED_PORTS as _allowed_ports_base,
    ALLOWED_SCHEMES as _allowed_schemes_base,
    ALLOWED_VIDEO_MIME_TYPES,
    MAX_URL_LENGTH,
    MAGIC_NUMBERS,
    UrlSecurityError as _UrlSecurityError_base,
    check_internal_hostname as _check_internal_hostname_base,
    check_ssrf_ip as _check_ssrf_ip_base,
    is_ip_address as _is_ip_address_base,
    is_trusted_domain as _is_trusted_domain_base,
    validate_magic_number as _validate_magic_number_base,
    validate_url_basic as _validate_url_basic_base,
)

logger = logging.getLogger(__name__)

# ── 兼容导出（保持原有变量名供外部引用） ──────────────────────────────────
ALLOWED_SCHEMES = set(_allowed_schemes_base)
ALLOWED_PORTS = set(_allowed_ports_base)
UrlSecurityError = _UrlSecurityError_base

# 可信域名白名单（从环境变量读取）
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 直接访问
ALLOW_DIRECT_IP = os.environ.get("URL_SECURITY_ALLOW_DIRECT_IP", "false").lower() == "true"

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

# 下载块大小
_DOWNLOAD_CHUNK_SIZE = 8192

# 最大重定向次数
_MAX_REDIRECTS = 5

# 魔数校验最大读取字节数
_MAGIC_CHECK_READ_SIZE = 256


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

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


# ── DNS 解析 SSRF 检查（有副作用） ─────────────────────────────────────────


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:
                _check_ssrf_ip_base(ip_str)
            except ValueError:
                continue
    except socket.gaierror as e:
        raise UrlSecurityError(f"域名解析失败: {hostname} ({e})") from e


def _validate_magic_number(file_path: str, allowed_mime_types: set[str]) -> None:
    """校验文件头魔数（从文件读取后委托给 domain 纯逻辑）."""
    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

    _validate_magic_number_base(header, allowed_mime_types)


# ── 对外 API ────────────────────────────────────────────────────────────────


def validate_url_safety(url: str, *, purpose: str = "download") -> str:
    """校验 URL 安全性，返回标准化后的 URL（含 DNS 解析 SSRF 检查）.

    Args:
        url: 待校验的 URL
        purpose: 用途描述（用于日志）

    Returns:
        标准化后的 URL

    Raises:
        UrlSecurityError: URL 不安全
    """
    # 基础校验（纯逻辑，不含 DNS）
    _validate_url_basic_base(
        url,
        trusted_domains=TRUSTED_DOMAINS,
        allow_direct_ip=ALLOW_DIRECT_IP,
    )

    # 如果 hostname 是域名（不是 IP），做 DNS 解析 SSRF 检查
    from urllib.parse import urlparse

    parsed = urlparse(url)
    hostname = parsed.hostname
    if hostname and not _is_ip_address_base(hostname):
        try:
            _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

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


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 + 每次重定向后都校验）
    - 重定向次数限制 + 手动跟随
    - 文件大小限制（流式读取）
    - MIME 类型白名单 + 魔数二次校验

    Returns:
        实际下载的字节数
    """
    current_url = url
    redirect_count = 0
    total_bytes = 0

    no_redirect_opener = urllib.request.build_opener(NoRedirectHandler())

    while True:
        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:
            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:
            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}, 允许: {sorted(allowed_mime_types)}")

            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)

            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 并返回字节内容（适合小文件）。"""
    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