from __future__ import annotations import fnmatch import ipaddress import socket from typing import Any from urllib.parse import urlparse from yuxi.utils.logging_config import logger _UNSAFE_TRANSPORT_FIELDS = frozenset( { "agent", "cert", "cert_file", "key", "key_file", "dispatcher", "proxy", "session", } ) _AUTH_RESPONSE_MAX_BYTES = 2 * 1024 * 1024 _PRIVATE_IP_RANGES = [ ipaddress.ip_network("10.0.0.0/8"), ipaddress.ip_network("172.16.0.0/12"), ipaddress.ip_network("192.168.0.0/16"), ipaddress.ip_network("127.0.0.0/8"), ipaddress.ip_network("169.254.0.0/16"), ipaddress.ip_network("0.0.0.0/8"), ipaddress.ip_network("224.0.0.0/4"), ipaddress.ip_network("fc00::/7"), ipaddress.ip_network("fe80::/10"), ipaddress.ip_network("::1/128"), ipaddress.ip_network("::/128"), ] class NetworkGuardError(ValueError): pass def is_private_url(url: str) -> bool: if not url: return False try: parsed = urlparse(url) hostname = parsed.hostname if not hostname: return False if hostname in ("localhost", "0.0.0.0", "[::]", "127.0.0.1", "::1"): return True addr = ipaddress.ip_address(hostname) return any(addr in net for net in _PRIVATE_IP_RANGES) except ValueError: try: resolved = socket.getaddrinfo(hostname, None) for _, _, _, _, sockaddr in resolved: ip = sockaddr[0] addr = ipaddress.ip_address(ip) if any(addr in net for net in _PRIVATE_IP_RANGES): return True except (socket.gaierror, OSError): return False return False def sanitize_transport_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: cleaned: dict[str, Any] = {} removed: list[str] = [] for key, value in kwargs.items(): if key in _UNSAFE_TRANSPORT_FIELDS: removed.append(key) continue cleaned[key] = value if removed: logger.warning(f"[SSRF] Removed unsafe transport fields: {removed}") return cleaned def check_response_size(data: bytes) -> bytes: if len(data) > _AUTH_RESPONSE_MAX_BYTES: logger.warning( f"[SSRF] Response exceeds size limit: {len(data)} bytes > {_AUTH_RESPONSE_MAX_BYTES} bytes, truncating" ) return data[:_AUTH_RESPONSE_MAX_BYTES] return data def check_response_text(text: str) -> str: data = text.encode("utf-8", errors="replace") truncated = check_response_size(data) return truncated.decode("utf-8", errors="replace") def validate_url(url: str, allow_private: bool = False) -> str: if not allow_private and is_private_url(url): raise ValueError(f"URL {url} resolves to a private/internal address") return url def safe_webhook_url(url: str) -> str: return validate_url(url, allow_private=False) def is_hostname_allowed(hostname: str, allowed_hosts: list[str]) -> bool: if not allowed_hosts: return False for pattern in allowed_hosts: if fnmatch.fnmatch(hostname, pattern): return True return False def validate_url_with_whitelist( url: str, allowed_hosts: list[str] | None = None, enforce_https: bool = True, allow_private: bool = False, ) -> str: parsed = urlparse(url) if enforce_https and parsed.scheme != "https": raise NetworkGuardError(f"URL scheme must be https, got: {parsed.scheme}") hostname = parsed.hostname if not hostname: raise NetworkGuardError(f"URL has no resolvable hostname: {url}") if not allow_private and is_private_url(url): raise NetworkGuardError(f"URL {url} resolves to a private/internal address") if allowed_hosts and not is_hostname_allowed(hostname, allowed_hosts): raise NetworkGuardError(f"Hostname '{hostname}' not in allowed hosts whitelist") return url def build_ssrf_safe_headers() -> dict[str, str]: return { "X-SSRF-Protection": "1", } async def fetch_with_ssrf_guard( url: str, *, method: str = "GET", headers: dict[str, str] | None = None, allowed_hosts: list[str] | None = None, enforce_https: bool = True, allow_private: bool = False, timeout: float = 30.0, **kwargs: Any, ) -> tuple[int, bytes]: import aiohttp validate_url_with_whitelist( url, allowed_hosts=allowed_hosts, enforce_https=enforce_https, allow_private=allow_private, ) merged_headers = build_ssrf_safe_headers() if headers: merged_headers.update(headers) safe_kwargs = sanitize_transport_kwargs(kwargs) try: async with aiohttp.ClientSession() as session: async with session.request( method, url, headers=merged_headers, timeout=aiohttp.ClientTimeout(total=timeout), **safe_kwargs ) as resp: data = await resp.read() return resp.status, check_response_size(data) except aiohttp.ClientError as e: raise NetworkGuardError(f"SSRF-guarded request failed for {url}: {e}") from e def apply_ssrf_guard_defaults() -> dict[str, Any]: return { "agent": None, "cert": None, "cert_file": None, "key": None, "key_file": None, }