Kris
9ce17c1662
- 新增完整的 SSRF 防护工具类与配置项 - 重构 URL 校验逻辑，替换原有白名单校验实现 - 新增 DNS 缓存与防重绑定防护 - 兼容原有 YUXI_URL_WHITELIST 环境变量配置 - 新增配套单元测试覆盖防护逻辑
118 lines
4.4 KiB
Python
118 lines
4.4 KiB
Python
from urllib.parse import urljoin, urlparse
|
|
|
|
import httpx
|
|
|
|
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
|
|
|
|
|
|
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 (SSRF protection, size limit, content type).
|
|
|
|
Args:
|
|
url: The URL to fetch.
|
|
max_size: Maximum allowed size in bytes.
|
|
|
|
Returns:
|
|
tuple: (content_bytes, final_url)
|
|
|
|
Raises:
|
|
ValueError: If validation fails or download error occurs.
|
|
"""
|
|
if not is_url_parsing_enabled():
|
|
raise ValueError("URL parsing feature is disabled")
|
|
|
|
validator = _get_validator()
|
|
|
|
# 初始 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
|
|
|
|
# 手动处理重定向，每次重定向重新校验
|
|
try:
|
|
async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client:
|
|
while True:
|
|
logger.info(f"Fetching URL: {current_url}")
|
|
|
|
headers = {
|
|
"User-Agent": (
|
|
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
|
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
|
"Chrome/91.0.4472.124 "
|
|
"Safari/537.36"
|
|
)
|
|
}
|
|
|
|
async with client.stream("GET", current_url, headers=headers) as response:
|
|
# Handle Redirects
|
|
if response.status_code in (301, 302, 303, 307, 308):
|
|
if redirect_count >= max_redirects:
|
|
raise ValueError("Too many redirects")
|
|
|
|
redirect_count += 1
|
|
location = response.headers.get("Location")
|
|
if not location:
|
|
raise ValueError("Redirect response missing Location header")
|
|
|
|
current_url = urljoin(current_url, location)
|
|
|
|
# 重定向 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
|
|
|
|
response.raise_for_status()
|
|
|
|
# Check Content-Type
|
|
content_type = response.headers.get("Content-Type", "").lower()
|
|
if not any(allowed in content_type for allowed in ALLOWED_CONTENT_TYPES):
|
|
raise ValueError(f"Unsupported Content-Type: {content_type}. Only HTML is supported.")
|
|
|
|
# Download content with size limit
|
|
content = bytearray()
|
|
async for chunk in response.aiter_bytes():
|
|
content.extend(chunk)
|
|
if len(content) > max_size:
|
|
raise ValueError(f"Content size exceeds limit of {max_size} bytes")
|
|
|
|
return bytes(content), current_url
|
|
|
|
except httpx.HTTPError as e:
|
|
logger.error(f"HTTP error fetching {url}: {e}")
|
|
raise ValueError(f"Failed to fetch URL: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error fetching {url}: {e}")
|
|
raise ValueError(f"Error fetching URL: {str(e)}")
|