Modules/_safe_http.py

"""SSRF-hardened HTTP GET for Web_Fetch.



Defence in depth:

1. URL check before every request and every redirect hop: http/https only,

   no credentials in the URL, allowed ports only, no internal host names,

   and every address the host resolves to must be public.

2. Connection-time check: the urllib3 connection re-resolves the host, rejects

   any non-public address and connects to the vetted IP literal. That closes

   the DNS-rebinding gap between step 1 and the actual connect.

3. Redirects are followed by hand (max 5) so each hop passes step 1.

4. Environment proxies are ignored (trust_env=False), so the connect really

   goes to the vetted address.

5. Body size (5 MB, measured after decompression) and total wall time (15 s)

   are capped.

"""

from __future__ import annotations

import ipaddress
import socket
import time
from dataclasses import dataclass
from urllib.parse import urljoin, urlsplit

import requests
from requests.adapters import HTTPAdapter
from urllib3.connection import HTTPConnection, HTTPSConnection
from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool
from urllib3.exceptions import NewConnectionError
from urllib3.util import connection as urllib3_connection

MAX_BYTES = 5 * 1024 * 1024
TOTAL_TIMEOUT_SECONDS = 15.0
CONNECT_TIMEOUT_SECONDS = 6.0
READ_TIMEOUT_SECONDS = 10.0
MAX_REDIRECTS = 5
ALLOWED_PORTS = {80, 443, 8080, 8443}

BLOCKED_HOST_SUFFIXES = (
    ".localhost",
    ".internal",
    ".local",
    ".localdomain",
    ".home.arpa",
    ".intranet",
    ".corp",
    ".lan",
)
BLOCKED_HOSTS = {"localhost", "localhost.localdomain", "ip6-localhost", "ip6-loopback", "metadata", "metadata.google.internal"}

_BLOCK_TAG = "[blocked-address]"

HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) "
        "Chrome/128.0.0.0 Safari/537.36"
    ),
    "Accept-Language": "en-US,en;q=0.9",
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,text/plain;q=0.8,*/*;q=0.5",
    "Accept-Encoding": "gzip, deflate",
    "DNT": "1",
    "Upgrade-Insecure-Requests": "1",
}


class FetchBlocked(Exception):
    """The URL points somewhere this tool refuses to go."""


class FetchFailed(Exception):
    """The request failed for an ordinary network/HTTP reason."""


def is_public_ip(value: str) -> bool:
    try:
        ip = ipaddress.ip_address(value.split("%", 1)[0])
    except ValueError:
        return False
    if isinstance(ip, ipaddress.IPv6Address):
        embedded = ip.ipv4_mapped or (ip.sixtofour if ip.sixtofour else None)
        if ip.teredo:
            embedded = ip.teredo[1]
        if embedded is not None and not is_public_ip(str(embedded)):
            return False
    if ip.is_multicast or ip.is_unspecified or ip.is_loopback or ip.is_link_local:
        return False
    if ip.is_private or ip.is_reserved:
        return False
    return bool(ip.is_global)


def _resolve(host: str, port: int) -> list[str]:
    infos = socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
    return list(dict.fromkeys(info[4][0] for info in infos))


def check_url(url: str) -> str:
    """Return the normalised URL or raise FetchBlocked with a readable reason."""
    try:
        parts = urlsplit(url.strip())
    except ValueError as exc:
        raise FetchBlocked(f"Malformed URL ({exc}).") from exc
    scheme = (parts.scheme or "").lower()
    if scheme not in ("http", "https"):
        raise FetchBlocked(f"Only http:// and https:// URLs are supported (got '{scheme or 'no scheme'}').")
    if parts.username or parts.password:
        raise FetchBlocked("URLs containing credentials (user:pass@host) are not allowed.")
    host = (parts.hostname or "").rstrip(".").lower()
    if not host:
        raise FetchBlocked("URL has no host name.")
    try:
        port = parts.port or (443 if scheme == "https" else 80)
    except ValueError as exc:
        raise FetchBlocked("URL has an invalid port.") from exc
    if port not in ALLOWED_PORTS:
        raise FetchBlocked(f"Port {port} is not allowed (allowed: {', '.join(map(str, sorted(ALLOWED_PORTS)))}).")
    if host in BLOCKED_HOSTS or host.endswith(BLOCKED_HOST_SUFFIXES):
        raise FetchBlocked(f"Host '{host}' is a local or internal name and cannot be fetched.")
    try:
        addresses = _resolve(host, port)
    except (socket.gaierror, UnicodeError) as exc:
        raise FetchFailed(f"Could not resolve host '{host}'.") from exc
    if not addresses:
        raise FetchFailed(f"Could not resolve host '{host}'.")
    for addr in addresses:
        if not is_public_ip(addr):
            raise FetchBlocked(f"Host '{host}' resolves to a private, loopback, link-local or reserved address ({addr}).")
    return parts.geturl()


class _GuardedConnectionMixin:
    """Re-resolve at connect time, refuse non-public addresses, connect to the vetted IP."""

    def _new_conn(self):  # type: ignore[override]
        host = self._dns_host  # type: ignore[attr-defined]
        port = self.port  # type: ignore[attr-defined]
        try:
            addresses = _resolve(host, port)
        except socket.gaierror as exc:
            raise NewConnectionError(self, f"Name resolution failed for {host}: {exc}") from exc  # type: ignore[arg-type]
        bad = [a for a in addresses if not is_public_ip(a)]
        if bad or not addresses:
            raise NewConnectionError(self, f"{_BLOCK_TAG} {host} resolves to non-public address {bad[0] if bad else '?'}")  # type: ignore[arg-type]
        last_exc: Exception | None = None
        for addr in addresses:
            try:
                return urllib3_connection.create_connection(
                    (addr, port),
                    self.timeout,  # type: ignore[attr-defined]
                    source_address=self.source_address,  # type: ignore[attr-defined]
                    socket_options=self.socket_options,  # type: ignore[attr-defined]
                )
            except OSError as exc:
                last_exc = exc
        raise NewConnectionError(self, f"Failed to establish a new connection: {last_exc}")  # type: ignore[arg-type]


class _GuardedHTTPConnection(_GuardedConnectionMixin, HTTPConnection):
    pass


class _GuardedHTTPSConnection(_GuardedConnectionMixin, HTTPSConnection):
    pass


class _GuardedHTTPPool(HTTPConnectionPool):
    ConnectionCls = _GuardedHTTPConnection


class _GuardedHTTPSPool(HTTPSConnectionPool):
    ConnectionCls = _GuardedHTTPSConnection


class _GuardedAdapter(HTTPAdapter):
    def init_poolmanager(self, *args, **kwargs):  # type: ignore[override]
        super().init_poolmanager(*args, **kwargs)
        self.poolmanager.pool_classes_by_scheme = {"http": _GuardedHTTPPool, "https": _GuardedHTTPSPool}


def _new_session() -> requests.Session:
    session = requests.Session()
    session.trust_env = False  # never route through env proxies or read .netrc
    adapter = _GuardedAdapter(max_retries=0, pool_connections=4, pool_maxsize=8)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    session.headers.update(HEADERS)
    return session


@dataclass
class FetchResult:
    url: str
    status: int
    content_type: str
    declared_charset: str | None
    body: bytes
    size_capped: bool


def _status_error(status: int) -> str:
    if status == 403:
        return "Access forbidden (HTTP 403). The website may be blocking automated requests."
    if status == 404:
        return "Page not found (HTTP 404). Please check the URL."
    if status == 429:
        return "Rate limited by the website (HTTP 429). Try again in a few minutes."
    if status >= 500:
        return f"The website returned a server error (HTTP {status})."
    return f"HTTP error {status}."


def _is_textual(content_type: str) -> bool:
    ctype = content_type.split(";", 1)[0].strip().lower()
    return ctype.startswith("text/") or ctype.endswith(("+xml", "+json", "/xml", "/json", "/xhtml")) or "html" in ctype


def safe_get(url: str) -> FetchResult:
    """GET a public http(s) URL with manual, re-validated redirects and hard caps."""
    started = time.monotonic()
    current = check_url(url)
    session = _new_session()
    try:
        for _hop in range(MAX_REDIRECTS + 1):
            remaining = TOTAL_TIMEOUT_SECONDS - (time.monotonic() - started)
            if remaining <= 0:
                raise FetchFailed(f"Request timed out (over {TOTAL_TIMEOUT_SECONDS:.0f} s).")
            try:
                resp = session.get(
                    current,
                    allow_redirects=False,
                    stream=True,
                    timeout=(min(CONNECT_TIMEOUT_SECONDS, remaining), min(READ_TIMEOUT_SECONDS, remaining)),
                )
            except requests.exceptions.ConnectionError as exc:
                if _BLOCK_TAG in str(exc):
                    raise FetchBlocked("The host resolved to a private, loopback, link-local or reserved address.") from exc
                if isinstance(exc, requests.exceptions.ConnectTimeout):
                    raise FetchFailed("Connection timed out. The website took too long to respond.") from exc
                raise FetchFailed("Connection error. The site may be down or unreachable.") from exc
            except requests.exceptions.Timeout as exc:
                raise FetchFailed("Request timed out. The website took too long to respond.") from exc
            except requests.exceptions.RequestException as exc:
                raise FetchFailed(f"Request failed: {str(exc)[:200]}") from exc

            if resp.is_redirect or resp.status_code in (301, 302, 303, 307, 308):
                location = resp.headers.get("Location")
                resp.close()
                if not location:
                    raise FetchFailed(f"Redirect (HTTP {resp.status_code}) without a Location header.")
                current = check_url(urljoin(current, location))
                continue

            if resp.status_code >= 400:
                status = resp.status_code
                resp.close()
                raise FetchFailed(_status_error(status))

            content_type = resp.headers.get("Content-Type", "")
            if content_type and not _is_textual(content_type):
                # Do not download bodies (PDFs, images, archives) the tool cannot use.
                resp.close()
                return FetchResult(current, resp.status_code, content_type, None, b"", False)
            declared = requests.utils.get_encoding_from_headers({"content-type": content_type}) if "charset" in content_type.lower() else None
            chunks: list[bytes] = []
            total = 0
            capped = False
            try:
                for chunk in resp.iter_content(chunk_size=65536):
                    if not chunk:
                        continue
                    chunks.append(chunk)
                    total += len(chunk)
                    if total >= MAX_BYTES:
                        capped = True
                        break
                    if time.monotonic() - started > TOTAL_TIMEOUT_SECONDS:
                        raise FetchFailed(f"Download exceeded {TOTAL_TIMEOUT_SECONDS:.0f} s and was stopped.")
            except requests.exceptions.RequestException as exc:
                raise FetchFailed(f"Download failed: {str(exc)[:200]}") from exc
            finally:
                resp.close()
            body = b"".join(chunks)[:MAX_BYTES]
            return FetchResult(current, resp.status_code, content_type, declared, body, capped)
        raise FetchFailed(f"Too many redirects (more than {MAX_REDIRECTS}).")
    finally:
        session.close()


__all__ = ["FetchBlocked", "FetchFailed", "FetchResult", "safe_get", "check_url", "is_public_ip", "MAX_BYTES"]