"""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 as _allowed_audio_base from packages.domain.url_security import ALLOWED_IMAGE_MIME_TYPES as _allowed_image_base from packages.domain.url_security import ALLOWED_PORTS as _allowed_ports_base from packages.domain.url_security import ALLOWED_SCHEMES as _allowed_schemes_base from packages.domain.url_security import ALLOWED_VIDEO_MIME_TYPES as _allowed_video_base from packages.domain.url_security import MAX_URL_LENGTH as _max_url_length_base from packages.domain.url_security import UrlSecurityError as _UrlSecurityError_base from packages.domain.url_security import check_internal_hostname as _check_internal_hostname_base from packages.domain.url_security import check_ssrf_ip as _check_ssrf_ip_base from packages.domain.url_security import is_ip_address as _is_ip_address_base from packages.domain.url_security import is_trusted_domain as _is_trusted_domain_base from packages.domain.url_security import validate_magic_number as _validate_magic_number_base from packages.domain.url_security import validate_url_basic as _validate_url_basic_base logger = logging.getLogger(__name__) # ── 兼容导出（保持原有变量名供外部引用） ────────────────────────────────── ALLOWED_SCHEMES = set(_allowed_schemes_base) ALLOWED_PORTS = set(_allowed_ports_base) ALLOWED_AUDIO_MIME_TYPES = set(_allowed_audio_base) ALLOWED_IMAGE_MIME_TYPES = set(_allowed_image_base) ALLOWED_VIDEO_MIME_TYPES = set(_allowed_video_base) MAX_URL_LENGTH = _max_url_length_base UrlSecurityError = _UrlSecurityError_base # 私有别名（供测试和内部引用） _check_internal_hostnames = _check_internal_hostname_base _check_ssrf_ip = _check_ssrf_ip_base def _is_trusted_domain(hostname: str) -> bool: """便捷包装：使用模块级 TRUSTED_DOMAINS 做可信域名检查.""" return _is_trusted_domain_base(hostname, TRUSTED_DOMAINS) # 可信域名白名单（从环境变量读取） 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 = 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 _is_ip_address_base(hostname): # IP 直接访问：通过本地别名调用以便 mock if ALLOW_DIRECT_IP: _check_ssrf_ip(hostname) elif 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