feat: 新增通用 SSRF 防护模块并重构 URL 校验逻辑
- 新增完整的 SSRF 防护工具类与配置项 - 重构 URL 校验逻辑，替换原有白名单校验实现 - 新增 DNS 缓存与防重绑定防护 - 兼容原有 YUXI_URL_WHITELIST 环境变量配置 - 新增配套单元测试覆盖防护逻辑
This commit is contained in:
parent
626d3da203
commit
9ce17c1662
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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,
|
||||
|
||||
312
backend/package/yuxi/utils/net_security.py
Normal file
312
backend/package/yuxi/utils/net_security.py
Normal file
@ -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()
|
||||
352
backend/test/unit/utils/test_net_security.py
Normal file
352
backend/test/unit/utils/test_net_security.py
Normal file
@ -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"
|
||||
Loading…
Reference in New Issue
Block a user