from __future__ import annotations import asyncio import ipaddress import logging import re import socket from typing import Any logger = logging.getLogger(__name__) _HOSTNAME_PATTERN = re.compile( r"^([a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)*[a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?$" ) _BUILTIN_SAFE_HOSTS: tuple[str, ...] = ( "dns.google", "one.one.one.one", "doh.opendns.com", "dns.quad9.net", ) # RFC 6761 special-use domain names that should be treated as internal _SPECIAL_USE_DOMAINS = frozenset( { "localhost", "local", "test", "invalid", "example", "localhost.localdomain", } ) _DEFAULT_DENY_BLOCKED_MARKER = "blocked by default-deny" class SsrfCheckResult: def __init__(self, safe: bool, reason: str = ""): self.safe = safe self.reason = reason def __repr__(self) -> str: return f"SsrfCheckResult(safe={self.safe}, reason={self.reason!r})" class SsrfGuard: def __init__(self, *, default_deny: bool = True): self._hostname_allowlist: list[str] = list(_BUILTIN_SAFE_HOSTS) self._pinned_hosts: dict[str, str] = {} self._default_deny = default_deny @property def default_deny(self) -> bool: return self._default_deny @property def allowlist(self) -> list[str]: return list(self._hostname_allowlist) def set_allowlist(self, patterns: list[str]) -> None: self._hostname_allowlist = patterns def pin_host(self, hostname: str, ip_address: str) -> None: self._pinned_hosts[hostname] = ip_address def check_hostname(self, hostname: str) -> SsrfCheckResult: if not self._is_valid_hostname(hostname): return SsrfCheckResult(safe=False, reason=f"Invalid hostname: {hostname}") if self._is_special_use_domain(hostname): return SsrfCheckResult(safe=False, reason=f"Special-use domain blocked: {hostname}") if hostname in self._pinned_hosts: return SsrfCheckResult(safe=True, reason="pinned host") if self._is_allowed_hostname(hostname): return SsrfCheckResult(safe=False, reason="allowlist matched, DNS resolution and IP check required") if self._default_deny: return SsrfCheckResult(safe=False, reason=_DEFAULT_DENY_BLOCKED_MARKER) return SsrfCheckResult(safe=False, reason="hostname requires DNS resolution and IP check") def check_ip(self, ip_str: str) -> SsrfCheckResult: try: ip = ipaddress.ip_address(ip_str) except ValueError: return SsrfCheckResult(safe=False, reason=f"Invalid IP address: {ip_str}") if self._is_private_ip(ip): return SsrfCheckResult(safe=False, reason=f"Private/internal IP blocked: {ip_str}") return SsrfCheckResult(safe=True, reason="public IP") def check_pinned_lookup(self, hostname: str, resolved_ip: str) -> SsrfCheckResult: pinned = self._pinned_hosts.get(hostname) if pinned is None: return SsrfCheckResult(safe=True, reason="no pinned lookup configured") if pinned != resolved_ip: return SsrfCheckResult( safe=False, reason=f"DNS pinned mismatch for {hostname}: expected {pinned}, got {resolved_ip}", ) return SsrfCheckResult(safe=True, reason="pinned lookup matched") async def check_url(self, hostname: str) -> SsrfCheckResult: result = self.check_hostname(hostname) if result.safe: return result if result.reason == _DEFAULT_DENY_BLOCKED_MARKER: return SsrfCheckResult(safe=False, reason=f"hostname not in allowlist: {hostname}") try: ip_str = await self._resolve_hostname(hostname) except Exception: return SsrfCheckResult(safe=False, reason=f"DNS resolution failed for: {hostname}") result = self.check_ip(ip_str) if not result.safe: return result return self.check_pinned_lookup(hostname, ip_str) async def check_url_safe(self, hostname: str) -> bool: result = await self.check_url(hostname) return result.safe @classmethod def from_config(cls, config: dict[str, Any]) -> SsrfGuard: ssrf_config = config.get("security", {}).get("ssrf", {}) if not isinstance(ssrf_config, dict): ssrf_config = {} default_deny = ssrf_config.get("defaultDeny", True) if not isinstance(default_deny, bool): default_deny = True guard = cls(default_deny=default_deny) allowlist = ssrf_config.get("allowlist", []) if isinstance(allowlist, list) and allowlist: merged = list(_BUILTIN_SAFE_HOSTS) for entry in allowlist: if isinstance(entry, str) and entry.strip() and entry.strip() not in merged: merged.append(entry.strip()) guard._hostname_allowlist = merged pinned = ssrf_config.get("pinned", {}) if isinstance(pinned, dict): for host, ip in pinned.items(): if isinstance(host, str) and isinstance(ip, str): guard.pin_host(host.strip(), ip.strip()) return guard @staticmethod async def _resolve_hostname(hostname: str) -> str: loop = asyncio.get_running_loop() addrinfo = await loop.getaddrinfo(hostname, None, family=socket.AF_UNSPEC, type=socket.SOCK_STREAM) for family, _, _, _, sockaddr in addrinfo: if family in (socket.AF_INET, socket.AF_INET6): return sockaddr[0] raise OSError(f"No suitable address found for {hostname}") def _is_valid_hostname(self, hostname: str) -> bool: if len(hostname) > 253: return False if not _HOSTNAME_PATTERN.match(hostname): return False return True def _is_special_use_domain(self, hostname: str) -> bool: hostname_lower = hostname.lower() if hostname_lower in _SPECIAL_USE_DOMAINS: return True for domain in _SPECIAL_USE_DOMAINS: if hostname_lower.endswith(f".{domain}"): return True return False def _is_allowed_hostname(self, hostname: str) -> bool: if not self._hostname_allowlist: return False hostname_lower = hostname.lower() for pattern in self._hostname_allowlist: pattern_lower = pattern.lower() if pattern_lower.startswith("*."): suffix = pattern_lower[1:] if hostname_lower.endswith(suffix): return True elif pattern_lower == hostname_lower: return True return False @staticmethod def _is_private_ip(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: return ( ip.is_loopback or ip.is_private or ip.is_link_local or ip.is_multicast or ip.is_reserved or ip.is_unspecified ) ssrf_guard = SsrfGuard()