ForcePilot/backend/package/yuxi/channels/auth/ssrf_guard.py

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,
    }