Yash030/claude-code-proxy

1"""Egress policy for user-controlled web_fetch URLs (SSRF guard)."""2 3from __future__ import annotations4 5import ipaddress6import socket7from dataclasses import dataclass8from urllib.parse import urlparse9 10 11@dataclass(frozen=True, slots=True)12class WebFetchEgressPolicy:13    """Egress rules for user-influenced web_fetch URLs."""14 15    allow_private_network_targets: bool16    allowed_schemes: frozenset[str]17 18 19class WebFetchEgressViolation(ValueError):20    """Raised when a web_fetch URL is rejected by egress policy (SSRF guard)."""21 22 23def _port_for_url(parsed) -> int:24    if parsed.port is not None:25        return parsed.port26    return 443 if (parsed.scheme or "").lower() == "https" else 8027 28 29def _stream_getaddrinfo_or_raise(host: str, port: int) -> list[tuple]:30    try:31        return socket.getaddrinfo(32            host, port, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP33        )34    except OSError as exc:35        raise WebFetchEgressViolation(36            f"Could not resolve host {host!r}: {exc}"37        ) from exc38 39 40def get_validated_stream_addrinfos_for_egress(41    url: str, policy: WebFetchEgressPolicy42) -> list[tuple]:43    """Resolve and validate a URL for web_fetch, returning getaddrinfo rows for pinning.44 45    Each HTTP connect pins to only these `getaddrinfo` results so a malicious DNS46    server cannot rebind to a disallowed address between resolution and the TCP47    connect (used by :func:`api.web_tools.outbound._run_web_fetch`).48    """49    parsed = urlparse(url)50    scheme = (parsed.scheme or "").lower()51    if scheme not in policy.allowed_schemes:52        raise WebFetchEgressViolation(53            f"URL scheme {scheme!r} is not allowed for web_fetch"54        )55 56    host = parsed.hostname57    if host is None or host == "":58        raise WebFetchEgressViolation("web_fetch URL must include a host")59 60    port = _port_for_url(parsed)61 62    if policy.allow_private_network_targets:63        return _stream_getaddrinfo_or_raise(host, port)64 65    host_lower = host.lower()66    if host_lower == "localhost" or host_lower.endswith(".localhost"):67        raise WebFetchEgressViolation("localhost targets are not allowed for web_fetch")68    if host_lower.endswith(".local"):69        raise WebFetchEgressViolation(".local hostnames are not allowed for web_fetch")70 71    try:72        parsed_ip = ipaddress.ip_address(host)73    except ValueError:74        parsed_ip = None75 76    if parsed_ip is not None:77        if not parsed_ip.is_global:78            raise WebFetchEgressViolation(79                f"Non-public IP host {host!r} is not allowed for web_fetch"80            )81        return _stream_getaddrinfo_or_raise(host, port)82 83    infos = _stream_getaddrinfo_or_raise(host, port)84    for *_, sockaddr in infos:85        addr = sockaddr[0]86        try:87            resolved = ipaddress.ip_address(addr)88        except ValueError:89            continue90        if not resolved.is_global:91            raise WebFetchEgressViolation(92                f"Host {host!r} resolves to a non-public address ({resolved})"93            )94    return infos95 96 97def enforce_web_fetch_egress(url: str, policy: WebFetchEgressPolicy) -> None:98    """Validate ``url`` (scheme, host, and resolved addresses) for web_fetch."""99    get_validated_stream_addrinfos_for_egress(url, policy)100