wehub-resource-sync
131 lines
5.0 KiB
Python
131 lines
5.0 KiB
Python
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
|
|
|
|
# 最大允许下载大小 (例如 10MB)
|
|
MAX_DOWNLOAD_SIZE = 10 * 1024 * 1024
|
|
# 允许的 Content-Type
|
|
ALLOWED_CONTENT_TYPES = ["text/html", "application/xhtml+xml"]
|
|
|
|
|
|
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
|
|
|
|
|
|
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).
|
|
|
|
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")
|
|
|
|
# Initial validation
|
|
is_valid, error_msg = validate_url(url)
|
|
if not is_valid:
|
|
raise ValueError(f"Invalid URL: {error_msg}")
|
|
|
|
# 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")
|
|
|
|
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) "
|
|
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
|
"Chrome/91.0.4472.124 "
|
|
"Safari/537.36"
|
|
)
|
|
}
|
|
|
|
# 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):
|
|
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)
|
|
|
|
# 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")
|
|
|
|
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)}")
|