diff --git a/backend/package/yuxi/config/app.py b/backend/package/yuxi/config/app.py index 4b413852..8f372ed5 100644 --- a/backend/package/yuxi/config/app.py +++ b/backend/package/yuxi/config/app.py @@ -50,6 +50,11 @@ class Config(BaseModel): sandbox_max_output_bytes: int = Field(default=262144, description="沙箱最大输出字节数") sandbox_keepalive_interval_seconds: int = Field(default=30, description="沙箱保活间隔") + enable_ssrf_protection: bool = Field(default=True, description="是否启用 SSRF 防护") + ssrf_allowed_domains: list[str] = Field(default_factory=list, description="SSRF 域名白名单") + ssrf_blocked_domains: list[str] = Field(default_factory=list, description="SSRF 域名黑名单") + ssrf_dns_pinned_ttl: int = Field(default=300, description="SSRF DNS 缓存时间（秒）") + _config_file: Path | None = PrivateAttr(default=None) _user_modified_fields: set[str] = PrivateAttr(default_factory=set) @@ -111,6 +116,11 @@ class Config(BaseModel): if not self.sandbox_virtual_path_prefix.startswith("/"): self.sandbox_virtual_path_prefix = f"/{self.sandbox_virtual_path_prefix}" + # SSRF 环境变量覆盖（兼容 YUXI_URL_WHITELIST） + ssrf_env = os.getenv("YUXI_URL_WHITELIST", "") + if ssrf_env: + self.ssrf_allowed_domains = [d.strip() for d in ssrf_env.split(",") if d.strip()] + def save(self) -> None: if not self._config_file: logger.warning("Config file path not set") @@ -158,5 +168,15 @@ class Config(BaseModel): else: logger.warning(f"Unknown config key: {key}") + def get_ssrf_policy(self): + """根据配置生成 SSRF 策略。""" + from yuxi.utils.net_security import SSRFPolicy + + return SSRFPolicy( + allowed_domains=self.ssrf_allowed_domains, + blocked_domains=self.ssrf_blocked_domains, + dns_pinned_ttl=self.ssrf_dns_pinned_ttl, + ) + config = Config() diff --git a/backend/package/yuxi/knowledge/utils/url_fetcher.py b/backend/package/yuxi/knowledge/utils/url_fetcher.py index 7fa78e4f..5819e6ca 100644 --- a/backend/package/yuxi/knowledge/utils/url_fetcher.py +++ b/backend/package/yuxi/knowledge/utils/url_fetcher.py @@ -1,42 +1,34 @@ -import ipaddress -import socket from urllib.parse import urljoin, urlparse import httpx -from yuxi.knowledge.utils.url_validator import is_url_parsing_enabled, validate_url from yuxi.utils import logger +from yuxi.utils.net_security import SSRFValidator, SSRFPolicy, log_ssrf_block +from yuxi.knowledge.utils.url_validator import is_url_parsing_enabled # 最大允许下载大小 (例如 10MB) MAX_DOWNLOAD_SIZE = 10 * 1024 * 1024 # 允许的 Content-Type ALLOWED_CONTENT_TYPES = ["text/html", "application/xhtml+xml"] +_validator: SSRFValidator | None = None -async def is_private_ip(hostname: str) -> bool: - """Check if the hostname resolves to a private IP address.""" - import asyncio - try: - # Resolve hostname to IP in a separate thread to avoid blocking the event loop - ip_list = await asyncio.to_thread(socket.getaddrinfo, hostname, None) - for item in ip_list: - ip_addr = item[4][0] - ip_obj = ipaddress.ip_address(ip_addr) - if ip_obj.is_private or ip_obj.is_loopback or ip_obj.is_link_local: - return True - return False - except Exception as e: - logger.warning(f"Failed to resolve hostname {hostname}: {e}") - # If resolution fails, assume it's unsafe or let the connection fail naturally, - # but to be safe we can return True to block it if strict mode is preferred. - # For now, we return False assuming standard DNS failure handling. - return False +def _get_validator() -> SSRFValidator: + """获取 SSRF 验证器单例（从环境变量构建策略）。""" + global _validator + if _validator is None: + import os + + whitelist = [d.strip() for d in os.environ.get("YUXI_URL_WHITELIST", "").split(",") if d.strip()] + policy = SSRFPolicy(allowed_domains=whitelist) + _validator = SSRFValidator(policy) + return _validator async def fetch_url_content(url: str, max_size: int = MAX_DOWNLOAD_SIZE) -> tuple[bytes, str]: """ - Fetch URL content with security checks (size limit, content type, private IP blocking). + Fetch URL content with security checks (SSRF protection, size limit, content type). Args: url: The URL to fetch. @@ -51,27 +43,25 @@ async def fetch_url_content(url: str, max_size: int = MAX_DOWNLOAD_SIZE) -> tupl if not is_url_parsing_enabled(): raise ValueError("URL parsing feature is disabled") - # Initial validation - is_valid, error_msg = validate_url(url) - if not is_valid: - raise ValueError(f"Invalid URL: {error_msg}") + validator = _get_validator() - # Parse URL to check IP - parsed_url = urlparse(url) - if await is_private_ip(parsed_url.hostname): - raise ValueError("Access to private IP addresses is forbidden") + # 初始 SSRF 校验 + allowed, reason = await validator.validate_url(url) + if not allowed: + parsed = urlparse(url) + log_ssrf_block(parsed.hostname or url, reason, context="url-fetch") + raise ValueError(f"请求被安全策略阻止: {reason}") current_url = url redirect_count = 0 max_redirects = 5 - # We handle redirects manually to check each target URL against whitelist/private IP + # 手动处理重定向，每次重定向重新校验 try: async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client: while True: logger.info(f"Fetching URL: {current_url}") - # Request headers headers = { "User-Agent": ( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " @@ -81,7 +71,6 @@ async def fetch_url_content(url: str, max_size: int = MAX_DOWNLOAD_SIZE) -> tupl ) } - # Stream the response to check headers before downloading body async with client.stream("GET", current_url, headers=headers) as response: # Handle Redirects if response.status_code in (301, 302, 303, 307, 308): @@ -95,14 +84,12 @@ async def fetch_url_content(url: str, max_size: int = MAX_DOWNLOAD_SIZE) -> tupl current_url = urljoin(current_url, location) - # Validate the new URL - is_valid, error_msg = validate_url(current_url) - if not is_valid: - raise ValueError(f"Redirected to invalid URL: {error_msg}") - - parsed_new = urlparse(current_url) - if await is_private_ip(parsed_new.hostname): - raise ValueError("Redirected to private IP address") + # 重定向 SSRF 校验 + allowed, reason = await validator.validate_url(current_url) + if not allowed: + parsed = urlparse(current_url) + log_ssrf_block(parsed.hostname or current_url, reason, context="url-fetch-redirect") + raise ValueError(f"重定向被安全策略阻止: {reason}") continue # Start new request diff --git a/backend/package/yuxi/knowledge/utils/url_validator.py b/backend/package/yuxi/knowledge/utils/url_validator.py index e6109508..ff6d5098 100644 --- a/backend/package/yuxi/knowledge/utils/url_validator.py +++ b/backend/package/yuxi/knowledge/utils/url_validator.py @@ -1,24 +1,20 @@ -"""URL validation utilities for whitelist-based URL parsing.""" +"""URL validation utilities (backed by net_security.SSRFValidator).""" import os -from urllib.parse import urlparse -# Environment variable name for URL whitelist +from yuxi.utils.net_security import SSRFPolicy, SSRFValidator + YUXI_URL_WHITELIST_ENV = "YUXI_URL_WHITELIST" -def _get_whitelist() -> list[str]: - """Get the URL whitelist from environment variables.""" - whitelist_str = os.environ.get(YUXI_URL_WHITELIST_ENV, "") - if not whitelist_str: - return [] - # Split by comma and clean up whitespace - return [item.strip() for item in whitelist_str.split(",") if item.strip()] +def _get_policy_from_env() -> SSRFPolicy: + """从环境变量构建策略（兼容 YUXI_URL_WHITELIST）。""" + whitelist = [d.strip() for d in os.environ.get(YUXI_URL_WHITELIST_ENV, "").split(",") if d.strip()] + return SSRFPolicy(allowed_domains=whitelist) def validate_url(url: str) -> tuple[bool, str]: - """ - Validate if a URL is in the whitelist. + """Validate if a URL is allowed (sync wrapper). Args: url: The URL to validate. @@ -30,56 +26,25 @@ def validate_url(url: str) -> tuple[bool, str]: if not url: return False, "URL cannot be empty" - # Parse the URL - try: - parsed = urlparse(url) - except Exception as e: - return False, f"Invalid URL format: {e}" - - # Check if URL has a valid scheme - if not parsed.scheme: - return False, "URL must start with http:// or https://" - - if parsed.scheme not in ("http", "https"): - return False, "URL must use HTTP or HTTPS protocol" - - # Get the hostname - hostname = parsed.hostname - if not hostname: - return False, "Invalid URL: no hostname found" - - # Get the whitelist - whitelist = _get_whitelist() - - # If whitelist is empty, URL parsing is disabled - if not whitelist: + # 白名单为空时，URL 抓取功能禁用（保持原有行为） + if not is_url_parsing_enabled(): return False, "URL parsing feature is disabled" - # Check if hostname is in whitelist - for allowed in whitelist: - # Handle wildcard patterns like *.example.com - if allowed.startswith("*."): - domain = allowed[2:] - # Check if hostname ends with the domain (or matches exactly) - if hostname == domain or hostname.endswith(f".{domain}"): - return True, "" - else: - # Exact match or subdomain match - if hostname == allowed or hostname.endswith(f".{allowed}"): - return True, "" - - return False, f"Domain '{hostname}' is not in the whitelist" + validator = SSRFValidator(_get_policy_from_env()) + allowed, reason = validator.validate_url_sync(url) + if not allowed: + return False, reason + return True, "" def is_url_parsing_enabled() -> bool: """Check if URL parsing is enabled (whitelist is configured).""" - whitelist = _get_whitelist() - return len(whitelist) > 0 + return bool(os.environ.get(YUXI_URL_WHITELIST_ENV, "").strip()) def get_whitelist_info() -> dict: """Get information about the current whitelist configuration.""" - whitelist = _get_whitelist() + whitelist = [d.strip() for d in os.environ.get(YUXI_URL_WHITELIST_ENV, "").split(",") if d.strip()] return { "enabled": len(whitelist) > 0, "domains": whitelist, diff --git a/backend/package/yuxi/utils/net_security.py b/backend/package/yuxi/utils/net_security.py new file mode 100644 index 00000000..4167fbc3 --- /dev/null +++ b/backend/package/yuxi/utils/net_security.py @@ -0,0 +1,312 @@ +"""通用 SSRF 防护模块。 + +提供纵深防御的 URL 安全校验能力，防止 Agent 通过工具访问内网服务。 +校验顺序：协议校验 → 主机名阻止(Pre-DNS) → 域名白/黑名单(Pre-DNS) → IP 字面量校验(Pre-DNS) → DNS 解析后 IP 校验(Post-DNS) +""" + +from __future__ import annotations + +import ipaddress +import re +import socket +import time +from urllib.parse import urlparse + +from pydantic import BaseModel, Field, PrivateAttr + +from yuxi.utils.logging_config import logger + + +class SSRFPolicy(BaseModel): + """SSRF 防护策略配置。""" + + allowed_schemes: list[str] = Field( + default_factory=lambda: ["http", "https"], + description="允许的 URL 协议", + ) + blocked_ip_ranges: list[str] = Field( + default_factory=lambda: [ + "0.0.0.0/8", # 当前网络 + "10.0.0.0/8", # 私有网络 A 类 + "127.0.0.0/8", # 回环地址 + "169.254.0.0/16", # 链路本地（含云元数据 169.254.169.254） + "172.16.0.0/12", # 私有网络 B 类 + "192.168.0.0/16", # 私有网络 C 类 + "198.18.0.0/15", # RFC 2544 基准测试（含 fake-ip 代理） + "100.64.0.0/10", # Carrier-grade NAT + "224.0.0.0/4", # 组播 + "240.0.0.0/4", # 保留 + "255.255.255.255/32", # 广播 + "::/128", # IPv6 未指定 + "::1/128", # IPv6 回环 + "fc00::/7", # IPv6 唯一本地 + "fe80::/10", # IPv6 链路本地 + "ff00::/8", # IPv6 组播 + "100::/64", # IPv6 丢弃前缀 + "2001:2::/48", # IPv6 基准测试 + "2001:20::/28", # ORCHIDv2 + "fec0::/10", # 废弃的 site-local + ], + description="阻止的 IP 地址范围（CIDR）", + ) + allowed_domains: list[str] = Field( + default_factory=list, + description="域名白名单（空=不限制域名，仅限制 IP）。支持通配符 *.example.com", + ) + blocked_domains: list[str] = Field( + default_factory=list, + description="域名黑名单（精确匹配）", + ) + blocked_hostnames: list[str] = Field( + default_factory=lambda: [ + "localhost", + "localhost.localdomain", + "metadata.google.internal", + ], + description="阻止的主机名（精确匹配，Pre-DNS 阶段即阻止）", + ) + blocked_hostname_suffixes: list[str] = Field( + default_factory=lambda: [".localhost", ".local", ".internal"], + description="阻止的主机名后缀（Pre-DNS 阶段即阻止）", + ) + dns_pinned_ttl: int = Field( + default=300, + description="DNS 缓存时间（秒），防重绑定", + ) + + _blocked_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = PrivateAttr(default_factory=list) + _blocked_domains_normalized: set[str] = PrivateAttr(default_factory=set) + + model_config = {"arbitrary_types_allowed": True} + + def model_post_init(self, __context) -> None: + self._blocked_networks = [ + ipaddress.ip_network(cidr) for cidr in self.blocked_ip_ranges + ] + self._blocked_domains_normalized = { + _normalize_hostname(d) for d in self.blocked_domains + } + + +class SSRFBlockedError(Exception): + """SSRF 防护拦截异常。""" + + def __init__(self, message: str, hostname: str = ""): + super().__init__(message) + self.hostname = hostname + + +def _normalize_hostname(hostname: str) -> str: + """主机名规范化：小写化、去尾点、去 IPv6 方括号。""" + normalized = hostname.strip().lower().rstrip(".") + if normalized.startswith("[") and normalized.endswith("]"): + normalized = normalized[1:-1] + return normalized + + +def _is_blocked_hostname(hostname: str, policy: SSRFPolicy) -> bool: + """检查主机名是否在阻止列表中（Pre-DNS，零副作用）。""" + normalized = _normalize_hostname(hostname) + if not normalized: + return True # fail-closed + if normalized in policy.blocked_hostnames: + return True + return any(normalized.endswith(suffix) for suffix in policy.blocked_hostname_suffixes) + + +def _matches_domain_pattern(hostname: str, pattern: str) -> bool: + """检查主机名是否匹配域名模式。 + + *.example.com 匹配子域名但不匹配 apex 域（example.com 本身不匹配 *.example.com）。 + """ + normalized = _normalize_hostname(hostname) + pattern = _normalize_hostname(pattern) + if pattern.startswith("*."): + suffix = pattern[2:] + if not suffix or normalized == suffix: + return False + return normalized.endswith(f".{suffix}") + return normalized == pattern or normalized.endswith(f".{pattern}") + + +def _is_allowed_by_domain_whitelist(hostname: str, allowed_domains: list[str]) -> bool: + """检查主机名是否通过域名白名单（空白名单=不限制）。""" + if not allowed_domains: + return True + return any(_matches_domain_pattern(hostname, d) for d in allowed_domains) + + +def _looks_like_ipv4_literal(hostname: str) -> bool: + """检测主机名是否看起来像 IPv4 字面量（含非标准格式）。 + + Python 的 ipaddress 不接受八进制、十六进制、短格式等，但浏览器/HTTP 库可能解析。 + """ + parts = hostname.split(".") + if not parts or len(parts) > 4: + return False + if any(not p for p in parts): + return True # 空段，异常格式 + return all( + p.isdigit() or (p.lower().startswith("0x") and all(c in "0123456789abcdef" for c in p[2:])) + for p in parts + ) + + +def _is_blocked_ip(addr: ipaddress.IPv4Address | ipaddress.IPv6Address, policy: SSRFPolicy) -> bool: + """检查 IP 是否在阻止范围内。""" + return any(addr in net for net in policy._blocked_networks) + + +def _check_ip_literal(hostname: str, policy: SSRFPolicy) -> tuple[bool, str]: + """检查主机名是否为 IP 字面量（含非标准格式和嵌入式 IPv4-in-IPv6）。 + + 返回 (is_handled, reason) — is_handled=True 表示已判定，is_handled=False 表示不是 IP 字面量。 + """ + normalized = _normalize_hostname(hostname) + + # 1. 标准 IPv4/IPv6 地址 + try: + ip = ipaddress.ip_address(normalized) + if _is_blocked_ip(ip, policy): + return True, f"目标 IP 在阻止范围内: {normalized}" + # IPv4-mapped IPv6: 检查嵌入式 IPv4（如 ::ffff:127.0.0.1） + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped: + if _is_blocked_ip(ip.ipv4_mapped, policy): + return True, f"IPv4-mapped 地址解析到被阻止的 IP: {normalized} -> {ip.ipv4_mapped}" + return True, "" # 合法公网 IP，通过 + except ValueError: + pass + + # 2. 畸形 IPv6 字面量（含冒号但无法解析） — fail-closed + if ":" in normalized: + try: + ipaddress.ip_address(normalized) + except ValueError: + return True, f"畸形 IPv6 地址，按安全策略阻止: {normalized}" + + # 3. 非标准 IPv4 字面量（八进制、十六进制、短格式） — fail-closed + if _looks_like_ipv4_literal(normalized): + return True, f"非标准 IPv4 字面量，按安全策略阻止: {normalized}" + + return False, "" # 不是 IP 字面量，是域名 + + +def log_ssrf_block(hostname: str, reason: str, context: str = "url-fetch") -> None: + """记录 SSRF 拦截事件（安全审计日志）。 + + 仅记录主机名，不记录完整 URL，避免泄露敏感参数。 + """ + logger.warning(f"security: blocked URL fetch ({context}) targetHost={hostname} reason={reason}") + + +class SSRFValidator: + """SSRF 防护验证器（async 优先，与项目风格一致）。""" + + def __init__(self, policy: SSRFPolicy | None = None): + self.policy = policy or SSRFPolicy() + self._dns_cache: dict[str, tuple[list[str], float]] = {} + + async def validate_url(self, url: str) -> tuple[bool, str]: + """校验 URL 是否允许访问（async，DNS 解析不阻塞事件循环）。 + + 校验顺序（纵深防御，Pre-DNS → Post-DNS）： + 1. URL 解析 + 2. 协议校验 + 3. 主机名阻止列表（Pre-DNS，零副作用） + 4. 域名白名单/黑名单（Pre-DNS） + 5. IP 字面量校验（含非标准格式和嵌入式 IPv4-in-IPv6，Pre-DNS） + 6. DNS 解析后 IP 校验（Post-DNS，防重绑定） + + Returns: + (是否通过, 原因描述) + """ + # 1. 解析 URL + try: + parsed = urlparse(url) + except Exception: + return False, "URL 格式无效" + + # 2. 协议校验 + if parsed.scheme not in self.policy.allowed_schemes: + return False, f"不允许的协议: {parsed.scheme}，仅允许 {self.policy.allowed_schemes}" + + hostname = parsed.hostname + if not hostname: + return False, "URL 缺少主机名" + + # 3. 主机名阻止列表（Pre-DNS） + if _is_blocked_hostname(hostname, self.policy): + return False, f"主机名在阻止列表中: {hostname}" + + # 4. 域名白名单/黑名单（Pre-DNS） + if not _is_allowed_by_domain_whitelist(hostname, self.policy.allowed_domains): + return False, f"域名不在白名单中: {hostname}" + + if _normalize_hostname(hostname) in self.policy._blocked_domains_normalized: + return False, f"域名在黑名单中: {hostname}" + + # 5. IP 字面量校验（Pre-DNS） + is_handled, reason = _check_ip_literal(hostname, self.policy) + if is_handled: + if reason: + return False, reason + # 合法公网 IP 字面量，通过 + return True, "通过" + + # 6. DNS 解析后 IP 校验（Post-DNS，防重绑定） + resolved_ips = await self._resolve_with_pin(hostname) + if not resolved_ips: + # DNS 解析失败 — fail-closed + return False, f"DNS 解析失败，按安全策略阻止: {hostname}" + + for ip_str in resolved_ips: + try: + ip = ipaddress.ip_address(ip_str) + if _is_blocked_ip(ip, self.policy): + return False, f"DNS 解析到被阻止的 IP: {hostname} -> {ip_str}" + except ValueError: + # 解析出非标准 IP — fail-closed + return False, f"DNS 解析到非标准地址，按安全策略阻止: {hostname} -> {ip_str}" + + return True, "通过" + + def validate_url_sync(self, url: str) -> tuple[bool, str]: + """同步版本的 validate_url，用于非 async 上下文。DNS 解析会阻塞。""" + import asyncio + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + raise RuntimeError("validate_url_sync 不能在已运行的事件循环中调用，请使用 validate_url") + return asyncio.run(self.validate_url(url)) + + async def _resolve_with_pin(self, hostname: str) -> list[str]: + """DNS 解析并缓存，防止重绑定攻击。 + + 在 dns_pinned_ttl 秒内，同一主机名返回缓存的 IP 列表。 + """ + import asyncio + + now = time.time() + if hostname in self._dns_cache: + cached_ips, cached_at = self._dns_cache[hostname] + if now - cached_at < self.policy.dns_pinned_ttl: + return cached_ips + + try: + addr_infos = await asyncio.to_thread( + socket.getaddrinfo, hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM + ) + resolved = list({addr[4][0] for addr in addr_infos}) + except (socket.gaierror, OSError): + resolved = [] # fail-closed: 调用方应检查空列表 + + self._dns_cache[hostname] = (resolved, now) + return resolved + + def clear_dns_cache(self) -> None: + """清除 DNS 缓存（配置变更时调用）。""" + self._dns_cache.clear() diff --git a/backend/test/unit/utils/test_net_security.py b/backend/test/unit/utils/test_net_security.py new file mode 100644 index 00000000..40f50e14 --- /dev/null +++ b/backend/test/unit/utils/test_net_security.py @@ -0,0 +1,352 @@ +"""SSRF 防护模块单元测试。""" + +import ipaddress +import pytest +from unittest.mock import patch, AsyncMock + +from yuxi.utils.net_security import ( + SSRFPolicy, + SSRFValidator, + SSRFBlockedError, + _normalize_hostname, + _is_blocked_hostname, + _matches_domain_pattern, + _is_allowed_by_domain_whitelist, + _looks_like_ipv4_literal, + _check_ip_literal, +) + + +class TestNormalizeHostname: + def test_lowercase(self): + assert _normalize_hostname("Example.COM") == "example.com" + + def test_strip_trailing_dot(self): + assert _normalize_hostname("example.com.") == "example.com" + + def test_strip_ipv6_brackets(self): + assert _normalize_hostname("[::1]") == "::1" + + def test_empty(self): + assert _normalize_hostname("") == "" + + def test_whitespace(self): + assert _normalize_hostname(" example.com ") == "example.com" + + +class TestBlockedHostname: + def test_localhost(self): + assert _is_blocked_hostname("localhost", SSRFPolicy()) + + def test_localhost_localdomain(self): + assert _is_blocked_hostname("localhost.localdomain", SSRFPolicy()) + + def test_metadata_google_internal(self): + assert _is_blocked_hostname("metadata.google.internal", SSRFPolicy()) + + def test_dot_localhost_suffix(self): + assert _is_blocked_hostname("sub.localhost", SSRFPolicy()) + + def test_dot_local_suffix(self): + assert _is_blocked_hostname("printer.local", SSRFPolicy()) + + def test_dot_internal_suffix(self): + assert _is_blocked_hostname("service.internal", SSRFPolicy()) + + def test_normal_hostname_not_blocked(self): + assert not _is_blocked_hostname("example.com", SSRFPolicy()) + + def test_empty_hostname_blocked(self): + assert _is_blocked_hostname("", SSRFPolicy()) # fail-closed + + def test_case_insensitive(self): + assert _is_blocked_hostname("Localhost", SSRFPolicy()) + assert _is_blocked_hostname("SERVICE.INTERNAL", SSRFPolicy()) + + +class TestDomainPattern: + def test_exact_match(self): + assert _matches_domain_pattern("api.example.com", "api.example.com") + + def test_subdomain_match(self): + assert _matches_domain_pattern("sub.api.example.com", "api.example.com") + + def test_wildcard_match_subdomain(self): + assert _matches_domain_pattern("sub.example.com", "*.example.com") + + def test_wildcard_no_match_apex(self): + # *.example.com 不匹配 example.com 本身 + assert not _matches_domain_pattern("example.com", "*.example.com") + + def test_wildcard_match_deep_subdomain(self): + assert _matches_domain_pattern("a.b.example.com", "*.example.com") + + def test_case_insensitive(self): + assert _matches_domain_pattern("Sub.Example.COM", "*.example.com") + + +class TestDomainWhitelist: + def test_empty_whitelist_allows_all(self): + assert _is_allowed_by_domain_whitelist("any.com", []) + + def test_whitelist_match(self): + assert _is_allowed_by_domain_whitelist("api.example.com", ["api.example.com"]) + + def test_whitelist_no_match(self): + assert not _is_allowed_by_domain_whitelist("other.com", ["api.example.com"]) + + def test_whitelist_wildcard(self): + assert _is_allowed_by_domain_whitelist("sub.example.com", ["*.example.com"]) + + +class TestIpv4LiteralDetection: + def test_octal(self): + assert _looks_like_ipv4_literal("0177.0.0.1") + + def test_hex(self): + assert _looks_like_ipv4_literal("0x7f.0.0.1") + + def test_short_format(self): + assert _looks_like_ipv4_literal("127.1") + + def test_decimal_only(self): + assert _looks_like_ipv4_literal("2130706433") + + def test_normal_hostname_not_detected(self): + assert not _looks_like_ipv4_literal("example.com") + + def test_standard_ipv4_detected(self): + # 标准 IPv4 也看起来像字面量，但会被标准解析优先处理 + assert _looks_like_ipv4_literal("127.0.0.1") + + def test_too_many_parts(self): + assert not _looks_like_ipv4_literal("1.2.3.4.5") + + def test_empty_part(self): + assert _looks_like_ipv4_literal("1..3.4") + + +class TestIpLiteralCheck: + def test_blocked_private_ipv4(self): + handled, reason = _check_ip_literal("192.168.1.1", SSRFPolicy()) + assert handled + assert reason + + def test_blocked_loopback(self): + handled, reason = _check_ip_literal("127.0.0.1", SSRFPolicy()) + assert handled + assert reason + + def test_blocked_cloud_metadata(self): + handled, reason = _check_ip_literal("169.254.169.254", SSRFPolicy()) + assert handled + assert reason + + def test_blocked_ipv4_mapped_ipv6(self): + handled, reason = _check_ip_literal("::ffff:127.0.0.1", SSRFPolicy()) + assert handled + assert reason + + def test_blocked_ipv6_loopback(self): + handled, reason = _check_ip_literal("::1", SSRFPolicy()) + assert handled + assert reason + + def test_allowed_public_ipv4(self): + handled, reason = _check_ip_literal("8.8.8.8", SSRFPolicy()) + assert handled + assert not reason + + def test_nonstandard_ipv4_blocked(self): + handled, reason = _check_ip_literal("0177.0.0.1", SSRFPolicy()) + assert handled + assert reason # fail-closed + + def test_domain_not_handled(self): + handled, reason = _check_ip_literal("example.com", SSRFPolicy()) + assert not handled + + def test_blocked_10_network(self): + handled, reason = _check_ip_literal("10.0.0.1", SSRFPolicy()) + assert handled + assert reason + + def test_blocked_172_network(self): + handled, reason = _check_ip_literal("172.16.0.1", SSRFPolicy()) + assert handled + assert reason + + def test_blocked_zero_network(self): + handled, reason = _check_ip_literal("0.0.0.0", SSRFPolicy()) + assert handled + assert reason + + def test_allowed_public_ipv6(self): + handled, reason = _check_ip_literal("2001:4860:4860::8888", SSRFPolicy()) + assert handled + assert not reason + + +class TestSSRFValidator: + def setup_method(self): + self.validator = SSRFValidator() + + @pytest.mark.parametrize("url", [ + "http://127.0.0.1:8080/admin", + "http://192.168.1.1/config", + "http://10.0.0.1/internal", + "http://172.16.0.1/secret", + "http://169.254.169.254/latest/meta-data/", + "http://0.0.0.0/", + "ftp://example.com/file", + "http://metadata.google.internal/computeMetadata/v1/", + "http://service.internal/api", + "http://printer.local/ipp", + ]) + @pytest.mark.asyncio + async def test_blocked_urls(self, url): + allowed, reason = await self.validator.validate_url(url) + assert not allowed, f"应阻止 {url}，但通过了校验" + + @pytest.mark.asyncio + async def test_blocked_localhost(self): + allowed, reason = await self.validator.validate_url("http://localhost:5432/db") + assert not allowed + + @pytest.mark.parametrize("url", [ + "https://api.openai.com/v1/chat/completions", + "https://www.baidu.com/search", + "http://example.com/page", + ]) + @pytest.mark.asyncio + async def test_allowed_urls(self, url): + validator = SSRFValidator() + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=["1.2.3.4"]): + allowed, reason = await validator.validate_url(url) + assert allowed, f"应允许 {url}，但被阻止: {reason}" + + @pytest.mark.asyncio + async def test_domain_whitelist(self): + policy = SSRFPolicy(allowed_domains=["api.example.com", "trusted.org"]) + validator = SSRFValidator(policy) + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=["1.2.3.4"]): + allowed, _ = await validator.validate_url("https://api.example.com/endpoint") + assert allowed + + allowed, _ = await validator.validate_url("https://other.com/endpoint") + assert not allowed + + @pytest.mark.asyncio + async def test_domain_whitelist_wildcard(self): + policy = SSRFPolicy(allowed_domains=["*.example.com"]) + validator = SSRFValidator(policy) + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=["1.2.3.4"]): + allowed, _ = await validator.validate_url("https://sub.example.com/endpoint") + assert allowed + + # *.example.com 不匹配 apex 域 + allowed, _ = await validator.validate_url("https://example.com/endpoint") + assert not allowed + + @pytest.mark.asyncio + async def test_domain_blacklist(self): + policy = SSRFPolicy(blocked_domains=["evil.com"]) + validator = SSRFValidator(policy) + allowed, _ = await validator.validate_url("https://evil.com/attack") + assert not allowed + + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=["1.2.3.4"]): + allowed, _ = await validator.validate_url("https://safe.com/page") + assert allowed + + @pytest.mark.asyncio + async def test_dns_resolution_blocked_private(self): + """DNS 解析到私有 IP 时应阻止""" + validator = SSRFValidator() + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=["10.0.0.1"]): + allowed, reason = await validator.validate_url("https://looks-legit.com/api") + assert not allowed + assert "DNS 解析到被阻止的 IP" in reason + + @pytest.mark.asyncio + async def test_dns_resolution_fail_closed(self): + """DNS 解析失败时应阻止（fail-closed）""" + validator = SSRFValidator() + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=[]): + allowed, reason = await validator.validate_url("https://unresolvable.example.com/api") + assert not allowed + assert "DNS 解析失败" in reason + + @pytest.mark.asyncio + async def test_ipv4_mapped_ipv6_blocked(self): + """IPv4-mapped IPv6 地址应提取并校验嵌入式 IPv4""" + allowed, reason = await self.validator.validate_url("http://[::ffff:127.0.0.1]/admin") + assert not allowed + + @pytest.mark.asyncio + async def test_nonstandard_ipv4_literal_blocked(self): + """非标准 IPv4 字面量应被阻止（fail-closed）""" + handled, reason = _check_ip_literal("0177.0.0.1", SSRFPolicy()) + assert handled + assert reason + + @pytest.mark.asyncio + async def test_empty_url(self): + allowed, reason = await self.validator.validate_url("") + assert not allowed + + @pytest.mark.asyncio + async def test_no_hostname(self): + allowed, reason = await self.validator.validate_url("http:///path") + assert not allowed + + @pytest.mark.asyncio + async def test_invalid_scheme(self): + allowed, reason = await self.validator.validate_url("ftp://example.com/file") + assert not allowed + assert "协议" in reason + + @pytest.mark.asyncio + async def test_dns_pinning(self): + """DNS Pinning 在 TTL 内返回缓存结果""" + validator = SSRFValidator(SSRFPolicy(dns_pinned_ttl=60)) + # 首次解析 + with patch.object(validator, "_resolve_with_pin", new_callable=AsyncMock, return_value=["1.2.3.4"]) as mock: + await validator.validate_url("https://example.com/page") + assert mock.call_count == 1 + + +class TestSSRFPolicyModel: + def test_default_policy(self): + policy = SSRFPolicy() + assert policy.allowed_schemes == ["http", "https"] + assert policy.dns_pinned_ttl == 300 + assert len(policy._blocked_networks) > 0 + + def test_custom_policy(self): + policy = SSRFPolicy( + allowed_domains=["api.example.com"], + blocked_domains=["evil.com"], + dns_pinned_ttl=60, + ) + assert policy.allowed_domains == ["api.example.com"] + assert policy.blocked_domains == ["evil.com"] + assert policy.dns_pinned_ttl == 60 + + def test_blocked_networks_initialized(self): + policy = SSRFPolicy() + assert ipaddress.ip_network("10.0.0.0/8") in policy._blocked_networks + assert ipaddress.ip_network("127.0.0.0/8") in policy._blocked_networks + assert ipaddress.ip_network("::1/128") in policy._blocked_networks + + def test_blocked_networks_includes_cloud_metadata(self): + policy = SSRFPolicy() + # 验证 169.254.169.254 在某个阻止网络内 + assert any(ipaddress.ip_address("169.254.169.254") in net for net in policy._blocked_networks) + + +class TestSSRFBlockedError: + def test_error_message(self): + error = SSRFBlockedError("blocked", hostname="evil.com") + assert str(error) == "blocked" + assert error.hostname == "evil.com"