mcp-server/spark_executor/tools/fetch_url.py

# coding=utf-8
"""
@Time   :2026/7/9
@Author :tao.chen

Generic HTTP GET proxy for the agent. Lets the agent fetch URLs on the
cluster's network when it cannot reach those hosts directly. Security:
URL host must share >= 2 labels of suffix with the connection's yarn_rm_url
host; IP literals and non-HTTP schemes are rejected. Reuses the connection's
saved auth/SSL config so the agent doesn't need cluster credentials.
"""
import ipaddress
import fnmatch
from urllib.parse import urlparse

import httpx

from common.logging import logger
from spark_executor.core.yarn_client import YarnClientConfig
from spark_executor.models import FetchUrlResult
from spark_executor.tools.connections import store as conn_store

_MAX_BODY_BYTES = 1_000_000  # 1 MB cap on response body
_REQUEST_TIMEOUT_SECONDS = 30


def _host_suffix_overlap(host_a: str, host_b: str, min_labels: int = 2) -> bool:
    """Return True if host_a and host_b share at least min_labels suffix labels."""
    labels_a = host_a.lower().split(".")
    labels_b = host_b.lower().split(".")
    n = 0
    i, j = len(labels_a) - 1, len(labels_b) - 1
    while i >= 0 and j >= 0 and labels_a[i] == labels_b[j]:
        n += 1
        i -= 1
        j -= 1
    return n >= min_labels


def _host_matches_any_glob(host: str, patterns: list[str]) -> bool:
    """True if `host` matches any of the fnmatch glob patterns.

    fnmatch is case-sensitive on Linux (our deployment target). `*` in a
    pattern does NOT cross '.' boundaries, so 'ccam*' matches 'ccam50' but
    NOT 'ccam50.evil.com' (which is what we want — domain boundary matters
    for SSRF).
    """
    host_labels = host.split(".")
    for p in patterns:
        pat_labels = p.split(".")
        if len(pat_labels) != len(host_labels):
            continue
        if all(fnmatch.fnmatchcase(h, pat) for h, pat in zip(host_labels, pat_labels)):
            return True
    return False


def _validate_url_host(
    url: str,
    yarn_rm_url: str | None,
    allowed_hosts: list[str] | None = None,
) -> None:
    """Raise ValueError if the URL is not allowed to be fetched.

    Allowed if EITHER:
    - URL host matches one of `allowed_hosts` fnmatch globs (if provided)
    - URL host shares >= 2 labels of suffix with `yarn_rm_url` host
    """
    parsed = urlparse(url)
    if parsed.scheme not in ("http", "https"):
        raise ValueError(
            f"URL scheme must be http or https, got {parsed.scheme!r}"
        )
    host = parsed.hostname
    if not host:
        raise ValueError(f"URL has no host: {url!r}")
    # IP literal check
    try:
        ipaddress.ip_address(host)
        raise ValueError(
            f"URL host {host!r} is an IP literal — IP targets are not allowed. "
            f"Use a hostname on the cluster network."
        )
    except ValueError as e:
        if "IP literal" in str(e):
            raise
        # not an IP, continue
    # Glob allowlist — if any pattern matches, host is allowed regardless of suffix
    if allowed_hosts and _host_matches_any_glob(host, allowed_hosts):
        return
    if not yarn_rm_url:
        raise ValueError(
            f"Connection has no yarn_rm_url set, cannot validate URL host "
            f"(no allowed_url_hosts match either). Save a Connection with "
            f"yarn_rm_url or allowed_url_hosts set."
        )
    anchor = urlparse(yarn_rm_url).hostname
    if not anchor:
        raise ValueError(f"Connection's yarn_rm_url has no host: {yarn_rm_url!r}")
    if not _host_suffix_overlap(host, anchor, min_labels=2):
        allowed_hint = (
            f" Or set allowed_url_hosts=['<glob>'] on the connection to allow "
            f"this host (e.g. ['ccam*'] for ccam1-ccam99)."
            if not allowed_hosts
            else f" (no allowed_url_hosts pattern matched either)"
        )
        raise ValueError(
            f"URL host {host!r} does not share enough suffix with the "
            f"connection's yarn_rm_url host {anchor!r} (need >= 2 labels of "
            f"common suffix).{allowed_hint} Reject this fetch to prevent SSRF."
        )

def fetch_url(url: str, connection_name: str) -> FetchUrlResult:
    """Proxy an HTTP GET to url using the auth/SSL settings of connection_name."""
    logger.debug(f"fetch_url enter url={url} connection_name={connection_name}")
    conn = conn_store.get(connection_name)
    if conn is None:
        raise KeyError(f"Connection not found: {connection_name}")

    _validate_url_host(url, conn.yarn_rm_url, conn.allowed_url_hosts)

    config = YarnClientConfig.from_connection(conn)
    resp = httpx.get(
        url,
        auth=config.auth_for_httpx(),
        verify=config.verify_for_httpx(),
        timeout=_REQUEST_TIMEOUT_SECONDS,
        follow_redirects=True,
    )

    body = resp.text[:_MAX_BODY_BYTES]
    truncated = len(resp.text) > _MAX_BODY_BYTES

    logger.info(
        f"fetch_url ok url={url} connection_name={connection_name} "
        f"status_code={resp.status_code} "
        f"content_type={resp.headers.get('content-type', '')} "
        f"body_bytes={len(body)} truncated={truncated}"
    )

    return FetchUrlResult(
        url=url,
        status_code=resp.status_code,
        content_type=resp.headers.get("content-type", ""),
        body=body,
        truncated=truncated,
    )