Kris
9ce17c1662
- 新增完整的 SSRF 防护工具类与配置项 - 重构 URL 校验逻辑，替换原有白名单校验实现 - 新增 DNS 缓存与防重绑定防护 - 兼容原有 YUXI_URL_WHITELIST 环境变量配置 - 新增配套单元测试覆盖防护逻辑
53 lines
1.6 KiB
Python
53 lines
1.6 KiB
Python
"""URL validation utilities (backed by net_security.SSRFValidator)."""
|
||
|
||
import os
|
||
|
||
from yuxi.utils.net_security import SSRFPolicy, SSRFValidator
|
||
|
||
YUXI_URL_WHITELIST_ENV = "YUXI_URL_WHITELIST"
|
||
|
||
|
||
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 allowed (sync wrapper).
|
||
|
||
Args:
|
||
url: The URL to validate.
|
||
|
||
Returns:
|
||
A tuple of (is_valid, error_message).
|
||
If is_valid is True, error_message will be empty.
|
||
"""
|
||
if not url:
|
||
return False, "URL cannot be empty"
|
||
|
||
# 白名单为空时，URL 抓取功能禁用（保持原有行为）
|
||
if not is_url_parsing_enabled():
|
||
return False, "URL parsing feature is disabled"
|
||
|
||
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)."""
|
||
return bool(os.environ.get(YUXI_URL_WHITELIST_ENV, "").strip())
|
||
|
||
|
||
def get_whitelist_info() -> dict:
|
||
"""Get information about the current whitelist configuration."""
|
||
whitelist = [d.strip() for d in os.environ.get(YUXI_URL_WHITELIST_ENV, "").split(",") if d.strip()]
|
||
return {
|
||
"enabled": len(whitelist) > 0,
|
||
"domains": whitelist,
|
||
"count": len(whitelist),
|
||
}
|