diff --git a/src/plyngent/tools/__init__.py b/src/plyngent/tools/__init__.py index 29d9448..3a9b2c2 100644 --- a/src/plyngent/tools/__init__.py +++ b/src/plyngent/tools/__init__.py @@ -28,6 +28,11 @@ from .file import move_path as move_path from .file import read_file as read_file from .file import tree as tree from .file import write_file as write_file +from .net import NET_TOOLS as NET_TOOLS +from .net import clear_private_grants as clear_private_grants +from .net import fetch as fetch +from .net import grant_private_host as grant_private_host +from .net import set_fetch_policy_confirm_hook as set_fetch_policy_confirm_hook from .plugins import load_plugin_tools as load_plugin_tools from .process import PROCESS_TOOLS as PROCESS_TOOLS from .process import ask_into_pty as ask_into_pty diff --git a/src/plyngent/tools/catalog.py b/src/plyngent/tools/catalog.py index 186f0fe..725f00c 100644 --- a/src/plyngent/tools/catalog.py +++ b/src/plyngent/tools/catalog.py @@ -213,11 +213,19 @@ def _ensure_builtin_definitions_registered(catalog: ToolCatalog) -> None: """ from plyngent.tools.chat import CHAT_TOOLS from plyngent.tools.file import FILE_TOOLS + from plyngent.tools.net import NET_TOOLS from plyngent.tools.process import PROCESS_TOOLS from plyngent.tools.todo import TODO_TOOLS from plyngent.tools.vcs import VCS_TOOLS - for definition in (*FILE_TOOLS, *PROCESS_TOOLS, *VCS_TOOLS, *CHAT_TOOLS, *TODO_TOOLS): + for definition in ( + *FILE_TOOLS, + *PROCESS_TOOLS, + *VCS_TOOLS, + *CHAT_TOOLS, + *TODO_TOOLS, + *NET_TOOLS, + ): if catalog.get(definition.name) is None: catalog.register(definition, source=_BUILTIN_SOURCE) @@ -235,6 +243,7 @@ def register_builtin_tools(*, force: bool = False) -> ToolCatalog: # Group packages pull their leaf modules (FILE_TOOLS etc. still exported). import plyngent.tools.chat as _chat_tools import plyngent.tools.file as _file_tools + import plyngent.tools.net as _net_tools import plyngent.tools.process as _process_tools import plyngent.tools.temp_workspace as _temp_workspace_tools import plyngent.tools.todo as _todo_tools @@ -243,6 +252,7 @@ def register_builtin_tools(*, force: bool = False) -> ToolCatalog: _ = ( _chat_tools, _file_tools, + _net_tools, _process_tools, _temp_workspace_tools, _todo_tools, diff --git a/src/plyngent/tools/danger.py b/src/plyngent/tools/danger.py index b3abdfd..58cde2c 100644 --- a/src/plyngent/tools/danger.py +++ b/src/plyngent/tools/danger.py @@ -233,12 +233,26 @@ def _open_pty_reason(args: Mapping[str, object]) -> str | None: return _shell_or_dash_c_reason(argv, via="open_pty") +def _fetch_reason(args: Mapping[str, object]) -> str | None: + """Soft-confirm cleartext HTTP and mutating fetch methods (not private-host policy).""" + from plyngent.tools.net.fetch import fetch_soft_reason + + method_obj = args.get("method", "GET") + url_obj = args.get("url", "") + body_obj = args.get("body") + method = method_obj if isinstance(method_obj, str) else "GET" + url = url_obj if isinstance(url_obj, str) else "" + body = body_obj if isinstance(body_obj, str) else None + return fetch_soft_reason(method, url, body) + + def classify_danger(name: str, args: Mapping[str, object]) -> str | None: # noqa: PLR0911 """Return a short reason if ``name``/``args`` need user confirm, else ``None``. Hard denylists (paths/commands) still raise independently. This only covers soft confirms for mutating tools and risky shell/REPL launches (interactive shells and ``python -c`` / ``bash -c`` one-liners). + Private/loopback fetch targets use a separate policy grant (not YOLO). """ if name == "delete_path": return _delete_path_reason(args) @@ -254,4 +268,6 @@ def classify_danger(name: str, args: Mapping[str, object]) -> str | None: # noq return _run_command_batch_reason(args) if name == "open_pty": return _open_pty_reason(args) + if name == "fetch": + return _fetch_reason(args) return None diff --git a/src/plyngent/tools/net/__init__.py b/src/plyngent/tools/net/__init__.py new file mode 100644 index 0000000..9695e6e --- /dev/null +++ b/src/plyngent/tools/net/__init__.py @@ -0,0 +1,8 @@ +from .fetch import fetch as fetch +from .grants import clear_private_grants as clear_private_grants +from .grants import grant_private_host as grant_private_host +from .grants import set_fetch_policy_confirm_hook as set_fetch_policy_confirm_hook + +NET_TOOLS = [ + fetch, +] diff --git a/src/plyngent/tools/net/client.py b/src/plyngent/tools/net/client.py new file mode 100644 index 0000000..060c0f0 --- /dev/null +++ b/src/plyngent/tools/net/client.py @@ -0,0 +1,367 @@ +"""Async HTTP fetch helper (niquests) with manual redirects and body caps.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import niquests + +from plyngent.tools.net.grants import ensure_host_allowed +from plyngent.tools.net.policy import ( + DEFAULT_MAX_BYTES, + DEFAULT_MAX_REDIRECTS, + DEFAULT_TIMEOUT_SECONDS, + FetchPolicyError, + is_https_to_http_downgrade, + parse_fetch_url, + resolve_redirect_url, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + +_REDIRECT_STATUS = frozenset({301, 302, 303, 307, 308}) +_TEXT_CONTENT_HINTS = ( + "text/", + "application/json", + "application/xml", + "application/javascript", + "application/xhtml", + "application/x-www-form-urlencoded", + "application/graphql", + "application/problem+json", + "application/ld+json", + "+json", + "+xml", +) +_BINARY_CONTENT_HINTS = ("octet-stream", "image/", "audio/", "video/") + + +@dataclass(slots=True) +class FetchResult: + status: int + final_url: str + content_type: str + body_text: str + body_bytes: int + truncated: bool + redirects: int + method: str + security: str + warnings: list[str] = field(default_factory=list) + body_kind: str = "text" # text | binary | empty + + +def _content_type(headers: Mapping[str, str] | object | None) -> str: + if headers is None: + return "" + get = getattr(headers, "get", None) + if not callable(get): + return "" + value = get("Content-Type") or get("content-type") or "" + return str(value) + + +def _looks_text(content_type: str, sample: bytes) -> bool: + lower = content_type.lower() + if any(hint in lower for hint in _TEXT_CONTENT_HINTS): + return True + if not sample: + return True + if b"\x00" in sample[:512]: + return False + return not any(hint in lower for hint in _BINARY_CONTENT_HINTS) + + +def _decode_body(data: bytes, content_type: str) -> str: + charset = "utf-8" + lower = content_type.lower() + if "charset=" in lower: + part = lower.split("charset=", 1)[1] + charset = part.split(";")[0].strip().strip('"') or "utf-8" + try: + return data.decode(charset, errors="replace") + except LookupError: + return data.decode("utf-8", errors="replace") + + +def _security_label(*, original_scheme: str, final_scheme: str, warnings: list[str]) -> str: + if any("https-to-http" in w for w in warnings): + return "https-to-http-redirect" + if final_scheme == "http" or original_scheme == "http": + return "cleartext-http" + return "https" + + +def _location_header(resp: object) -> str | None: + resp_headers = getattr(resp, "headers", None) + if resp_headers is None: + return None + value = resp_headers.get("Location") or resp_headers.get("location") + return None if value is None else str(value) + + +def _apply_redirect_method(status: int, method: str, body: bytes | None) -> tuple[str, bytes | None]: + if status in {302, 303} and method != "GET": + return "GET", None + return method, body + + +def _note_redirect_security( + *, + previous_scheme: str, + next_url: str, + next_scheme: str, + hop: int, + allow_http_downgrade: bool, + warnings: list[str], +) -> None: + if is_https_to_http_downgrade(previous_scheme, next_scheme): + if not allow_http_downgrade: + msg = f"blocked HTTPS→HTTP redirect to {next_url} (set allow_http_downgrade=true to override)" + raise FetchPolicyError(msg) + warnings.append(f"https-to-http-redirect hop {hop}: {next_url}") + elif next_scheme == "http": + warnings.append(f"cleartext HTTP at hop {hop}: {next_url}") + + +def _body_payload(data: bytes, content_type: str) -> tuple[str, str]: + """Return (body_kind, body_text).""" + if not data: + return "empty", "" + if _looks_text(content_type, data): + return "text", _decode_body(data, content_type) + preview = data[:64].hex() + return ( + "binary", + f"(binary body omitted; content_type={content_type!r}; bytes={len(data)}; hex_prefix={preview})", + ) + + +def _empty_redirect_result( + *, + status: int, + current_url: str, + content_type: str, + redirects: int, + method: str, + original_scheme: str, + scheme: str, + warnings: list[str], +) -> FetchResult: + return FetchResult( + status=status, + final_url=current_url, + content_type=content_type, + body_text="", + body_bytes=0, + truncated=False, + redirects=redirects, + method=method, + security=_security_label( + original_scheme=original_scheme, + final_scheme=scheme, + warnings=warnings, + ), + warnings=warnings, + body_kind="empty", + ) + + +async def http_fetch( + *, + method: str, + url: str, + headers: Mapping[str, str], + body: bytes | None, + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + max_bytes: int = DEFAULT_MAX_BYTES, + max_redirects: int = DEFAULT_MAX_REDIRECTS, + follow_redirects: bool = True, + allow_http_downgrade: bool = False, +) -> FetchResult: + """Perform *method* on *url* with SSRF checks on each hop.""" + if timeout_seconds <= 0: + msg = "timeout_seconds must be > 0" + raise FetchPolicyError(msg) + if max_bytes < 1: + msg = "max_bytes must be >= 1" + raise FetchPolicyError(msg) + if max_redirects < 0: + msg = "max_redirects must be >= 0" + raise FetchPolicyError(msg) + + current_url = parse_fetch_url(url).url + original_scheme = parse_fetch_url(current_url).scheme + warnings: list[str] = [] + redirects = 0 + active_method = method + active_body = body + timeout = float(timeout_seconds) + + async with niquests.AsyncSession() as session: + while True: + await ensure_host_allowed(current_url) + parsed = parse_fetch_url(current_url) + try: + resp = await session.request( + active_method, + current_url, + headers=dict(headers), + data=active_body, + timeout=timeout, + allow_redirects=False, + stream=True, + ) + except niquests.RequestException as exc: + msg = f"HTTP request failed: {exc}" + raise FetchPolicyError(msg) from exc + + status = int(getattr(resp, "status_code", 0) or 0) + content_type = _content_type(getattr(resp, "headers", {}) or {}) + is_redirect = status in _REDIRECT_STATUS + + if follow_redirects and is_redirect and redirects < max_redirects: + location = _location_header(resp) + await _close_response(resp) + if not location: + return _empty_redirect_result( + status=status, + current_url=current_url, + content_type=content_type, + redirects=redirects, + method=active_method, + original_scheme=original_scheme, + scheme=parsed.scheme, + warnings=warnings, + ) + next_url = resolve_redirect_url(current_url, location) + next_parsed = parse_fetch_url(next_url) + _note_redirect_security( + previous_scheme=parsed.scheme, + next_url=next_url, + next_scheme=next_parsed.scheme, + hop=redirects + 1, + allow_http_downgrade=allow_http_downgrade, + warnings=warnings, + ) + redirects += 1 + current_url = next_url + active_method, active_body = _apply_redirect_method(status, active_method, active_body) + continue + + if follow_redirects and is_redirect and redirects >= max_redirects: + await _close_response(resp) + msg = f"too many redirects (max {max_redirects})" + raise FetchPolicyError(msg) + + data, truncated = await _read_capped(resp, max_bytes=max_bytes) + await _close_response(resp) + + final_scheme = parse_fetch_url(current_url).scheme + if final_scheme == "http": + warnings.append(f"cleartext HTTP final_url={current_url}") + body_kind, body_text = _body_payload(data, content_type) + return FetchResult( + status=status, + final_url=current_url, + content_type=content_type, + body_text=body_text, + body_bytes=len(data), + truncated=truncated, + redirects=redirects, + method=active_method, + security=_security_label( + original_scheme=original_scheme, + final_scheme=final_scheme, + warnings=warnings, + ), + warnings=list(dict.fromkeys(warnings)), + body_kind=body_kind, + ) + + +def _append_chunk( + chunks: list[bytes], + total: int, + data: bytes, + *, + max_bytes: int, +) -> tuple[int, bool]: + """Append *data* under *max_bytes*; return (new_total, truncated).""" + if total + len(data) > max_bytes: + need = max_bytes - total + if need > 0: + chunks.append(data[:need]) + total += need + return total, True + chunks.append(data) + return total + len(data), False + + +async def _consume_async_stream(stream: Any, *, max_bytes: int) -> tuple[bytes, bool]: + chunks: list[bytes] = [] + total = 0 + async for chunk in stream: + if not chunk: + continue + data = chunk if isinstance(chunk, bytes) else bytes(chunk) + total, truncated = _append_chunk(chunks, total, data, max_bytes=max_bytes) + if truncated: + return b"".join(chunks), True + return b"".join(chunks), False + + +def _consume_sync_stream(stream: Any, *, max_bytes: int) -> tuple[bytes, bool]: + chunks: list[bytes] = [] + total = 0 + for chunk in stream: + if not chunk: + continue + data = chunk if isinstance(chunk, bytes) else bytes(chunk) + total, truncated = _append_chunk(chunks, total, data, max_bytes=max_bytes) + if truncated: + return b"".join(chunks), True + return b"".join(chunks), False + + +async def _read_full_content(resp: object, *, max_bytes: int) -> tuple[bytes, bool]: + content = getattr(resp, "content", None) + if content is not None and hasattr(content, "__await__"): + raw = await content + data = raw if isinstance(raw, bytes) else bytes(raw) + elif isinstance(content, bytes): + data = content + else: + data = b"" + if len(data) > max_bytes: + return data[:max_bytes], True + return data, False + + +async def _read_capped(resp: object, *, max_bytes: int) -> tuple[bytes, bool]: + """Read at most *max_bytes* from a streamed response. + + niquests ``AsyncResponse.iter_content`` is async and returns an async + generator after await; ``content`` is an async property (coroutine). + """ + iter_content = getattr(resp, "iter_content", None) + if callable(iter_content): + stream_obj: Any = iter_content(chunk_size=65_536) + if hasattr(stream_obj, "__await__"): + stream_obj = await stream_obj + if hasattr(stream_obj, "__aiter__"): + return await _consume_async_stream(stream_obj, max_bytes=max_bytes) + if hasattr(stream_obj, "__iter__"): + return _consume_sync_stream(stream_obj, max_bytes=max_bytes) + return await _read_full_content(resp, max_bytes=max_bytes) + + +async def _close_response(resp: object) -> None: + close = getattr(resp, "close", None) + if close is None: + return + result: Any = close() + if hasattr(result, "__await__"): + await result diff --git a/src/plyngent/tools/net/fetch.py b/src/plyngent/tools/net/fetch.py new file mode 100644 index 0000000..7665b81 --- /dev/null +++ b/src/plyngent/tools/net/fetch.py @@ -0,0 +1,140 @@ +"""Model-facing ``fetch`` tool (HTTP GET/POST/PUT/DELETE).""" + +from __future__ import annotations + +from plyngent.agent import ToolTag, tool +from plyngent.tools.net.client import http_fetch +from plyngent.tools.net.policy import ( + DEFAULT_MAX_BODY_CHARS_IN, + DEFAULT_MAX_BYTES, + DEFAULT_MAX_CHARS, + DEFAULT_MAX_REDIRECTS, + DEFAULT_TIMEOUT_SECONDS, + FetchPolicyError, + normalize_method, + normalize_request_headers, + parse_fetch_url, + soft_confirm_reason, +) + + +def format_fetch_result( + *, + status: int, + final_url: str, + content_type: str, + body_text: str, + body_bytes: int, + truncated: bool, + redirects: int, + method: str, + security: str, + warnings: list[str], + body_kind: str, + max_chars: int, +) -> str: + text = body_text + char_truncated = False + if max_chars >= 1 and len(text) > max_chars: + omitted = len(text) - max_chars + text = text[:max_chars] + f"\n...[truncated {omitted} characters]" + char_truncated = True + warn_line = "; ".join(warnings) if warnings else "" + parts = [ + f"status={status}", + f"method={method}", + f"final_url={final_url}", + f"content_type={content_type}", + f"body_kind={body_kind}", + f"bytes={body_bytes}", + f"truncated={'true' if truncated or char_truncated else 'false'}", + f"redirects={redirects}", + f"security={security}", + ] + if warn_line: + parts.append(f"warnings={warn_line}") + parts.append("--- body ---") + parts.append(text) + return "\n".join(parts) + + +@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO | ToolTag.TRUSTABLE) +async def fetch( + url: str, + *, + method: str = "GET", + headers: dict[str, str] | None = None, + body: str | None = None, + user_agent: str | None = None, + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + max_bytes: int = DEFAULT_MAX_BYTES, + max_chars: int = DEFAULT_MAX_CHARS, + max_redirects: int = DEFAULT_MAX_REDIRECTS, + follow_redirects: bool = True, + allow_http_downgrade: bool = False, +) -> str: + """HTTP request (GET/POST/PUT/DELETE); return status, metadata, and truncated body. + + Use for public docs/APIs or (after human policy allow) local/LAN servers. + ``user_agent`` sets User-Agent when provided and takes precedence over a + User-Agent entry in ``headers``. A headers-only User-Agent is kept as-is. + When both omit UA, a small default is used. + + Private/loopback/link-local hosts require an explicit human policy grant + (not skipped by YOLO). HTTPS→HTTP redirects are blocked unless + ``allow_http_downgrade`` is true. + """ + try: + verb = normalize_method(method) + parsed = parse_fetch_url(url) + hdrs = normalize_request_headers(headers, user_agent=user_agent) + body_bytes: bytes | None = None + if body is not None: + if len(body) > DEFAULT_MAX_BODY_CHARS_IN: + return f"error: request body too large ({len(body)} chars; max {DEFAULT_MAX_BODY_CHARS_IN})" + body_bytes = body.encode("utf-8") + result = await http_fetch( + method=verb, + url=parsed.url, + headers=hdrs, + body=body_bytes, + timeout_seconds=timeout_seconds, + max_bytes=max_bytes, + max_redirects=max_redirects, + follow_redirects=follow_redirects, + allow_http_downgrade=allow_http_downgrade, + ) + except FetchPolicyError as exc: + return f"error: {exc}" + except Exception as exc: # noqa: BLE001 — tool surface returns error text + return f"error: fetch failed: {exc}" + + return format_fetch_result( + status=result.status, + final_url=result.final_url, + content_type=result.content_type, + body_text=result.body_text, + body_bytes=result.body_bytes, + truncated=result.truncated, + redirects=result.redirects, + method=result.method, + security=result.security, + warnings=result.warnings, + body_kind=result.body_kind, + max_chars=max_chars, + ) + + +# Re-export for danger classifier without importing client. +def fetch_soft_reason(method: str, url: str, body: str | None) -> str | None: + try: + verb = normalize_method(method) + parsed = parse_fetch_url(url) + except FetchPolicyError: + return None + return soft_confirm_reason( + method=verb, + url=parsed.url, + scheme=parsed.scheme, + body_present=bool(body), + ) diff --git a/src/plyngent/tools/net/grants.py b/src/plyngent/tools/net/grants.py new file mode 100644 index 0000000..57f10e2 --- /dev/null +++ b/src/plyngent/tools/net/grants.py @@ -0,0 +1,127 @@ +"""Instance-scoped private-fetch host grants (not YOLO-skippable).""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +from plyngent.tools.net.policy import FetchPolicyError, HostClass, assess_host, grant_key, parse_fetch_url +from plyngent.tools.workspace import WorkspaceError, require_bound_instance + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from plyngent.tools.context import InstanceState + +# (host, port, url, timeout_seconds) -> True to allow for this process/instance. +type FetchPolicyConfirmHook = Callable[[str, int, str, float], bool] + + +def get_private_grants(instance: InstanceState | None = None) -> set[str]: + """Return the live grant key set (host:port strings).""" + inst = instance if instance is not None else require_bound_instance() + bag_obj = inst.extras.get("fetch_private_grants") + if isinstance(bag_obj, set): + return cast("set[str]", bag_obj) + empty: set[str] = set() + inst.extras["fetch_private_grants"] = empty + return empty + + +def grant_private_host(host: str, port: int, *, instance: InstanceState | None = None) -> str: + """Record a private host:port grant; returns the grant key.""" + key = grant_key(host, port) + get_private_grants(instance).add(key) + return key + + +def clear_private_grants(*, instance: InstanceState | None = None) -> None: + get_private_grants(instance).clear() + + +def has_private_grant(host: str, port: int, *, instance: InstanceState | None = None) -> bool: + return grant_key(host, port) in get_private_grants(instance) + + +def get_fetch_policy_confirm_hook(instance: InstanceState | None = None) -> FetchPolicyConfirmHook | None: + inst = instance if instance is not None else require_bound_instance() + hook = inst.extras.get("fetch_policy_confirm_hook") + if hook is None or not callable(hook): + return None + + def _as_hook(host: str, port: int, url: str, timeout_seconds: float) -> bool: + return bool(hook(host, port, url, timeout_seconds)) + + return _as_hook + + +def set_fetch_policy_confirm_hook( + hook: FetchPolicyConfirmHook | None, + *, + instance: InstanceState | None = None, +) -> None: + inst = instance if instance is not None else require_bound_instance() + inst.extras["fetch_policy_confirm_hook"] = hook + + +async def ensure_host_allowed( + url: str, + *, + policy_timeout_seconds: float | None = None, + instance: InstanceState | None = None, +) -> None: + """Raise :class:`FetchPolicyError` if *url*'s host may not be contacted. + + Public hosts pass. Forbidden (metadata) always fail. Private/loopback + requires an instance grant or a successful policy confirm hook (never YOLO). + """ + try: + inst = instance if instance is not None else require_bound_instance() + except WorkspaceError as exc: + msg = f"instance state is not bound for fetch policy: {exc}" + raise FetchPolicyError(msg) from exc + + parsed = parse_fetch_url(url) + assessment = await assess_host(parsed.host, parsed.port) + if assessment.classification is HostClass.PUBLIC: + return + if assessment.classification is HostClass.FORBIDDEN: + msg = f"fetch blocked forbidden host {parsed.host!r} (addresses: {', '.join(assessment.addresses)})" + raise FetchPolicyError(msg) + + # PRIVATE + if has_private_grant(parsed.host, parsed.port, instance=inst): + return + + from plyngent.tools.workspace import get_policy_confirm_timeout + + timeout = ( + float(policy_timeout_seconds) if policy_timeout_seconds is not None else float(get_policy_confirm_timeout()) + ) + hook = get_fetch_policy_confirm_hook(inst) + if hook is None: + msg = ( + f"fetch blocked private/loopback host {parsed.host}:{parsed.port} " + f"({assessment.reason}; no policy confirm hook / non-interactive deny)" + ) + raise FetchPolicyError(msg) + try: + allowed = bool(hook(parsed.host, parsed.port, parsed.url, timeout)) + except Exception as exc: + msg = f"fetch blocked private/loopback host {parsed.host}:{parsed.port} (confirm failed: {exc})" + raise FetchPolicyError(msg) from exc + if not allowed: + msg = ( + f"fetch blocked private/loopback host {parsed.host}:{parsed.port} " + f"(user declined or timed out after {timeout:g}s; not skipped by YOLO)" + ) + raise FetchPolicyError(msg) + _ = grant_private_host(parsed.host, parsed.port, instance=inst) + + +def format_grant_preview(host: str, port: int, url: str) -> Sequence[str]: + return ( + f"host: {host}", + f"port: {port}", + f"url: {url}", + f"grant_key: {grant_key(host, port)}", + ) diff --git a/src/plyngent/tools/net/policy.py b/src/plyngent/tools/net/policy.py new file mode 100644 index 0000000..5a95aa3 --- /dev/null +++ b/src/plyngent/tools/net/policy.py @@ -0,0 +1,286 @@ +"""URL / host policy for the fetch tool (SSRF, private grants, cleartext).""" + +from __future__ import annotations + +import ipaddress +import socket +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING +from urllib.parse import urljoin, urlparse, urlunparse + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + +ALLOWED_METHODS: frozenset[str] = frozenset({"GET", "POST", "PUT", "DELETE"}) +ALLOWED_SCHEMES: frozenset[str] = frozenset({"http", "https"}) + +# Hop-by-hop / identity headers the model must not set (User-Agent is allowed). +_FORBIDDEN_REQUEST_HEADERS: frozenset[str] = frozenset( + { + "host", + "content-length", + "transfer-encoding", + "connection", + "keep-alive", + "upgrade", + "te", + "trailer", + "proxy-authorization", + "proxy-authenticate", + } +) + +DEFAULT_USER_AGENT = "plyngent-fetch/0.2" +DEFAULT_MAX_REDIRECTS = 5 +DEFAULT_TIMEOUT_SECONDS = 30.0 +DEFAULT_MAX_BYTES = 1_000_000 +DEFAULT_MAX_CHARS = 32_000 +DEFAULT_MAX_BODY_CHARS_IN = 256_000 # request body size limit (tool arg) + +# Cloud / link-local style targets never granted via policy UI. +_NEVER_ALLOW_NETWORKS: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = ( + ipaddress.ip_network("169.254.169.254/32"), # AWS/GCP-style metadata (IPv4) + ipaddress.ip_network("fd00:ec2::254/128"), # AWS IMDS IPv6 +) + + +class HostClass(Enum): + PUBLIC = "public" + PRIVATE = "private" # loopback, RFC1918, ULA, link-local, etc. (grantable) + FORBIDDEN = "forbidden" # metadata / never allow + + +class FetchPolicyError(ValueError): + """Hard fetch policy violation (returned to the model as error text).""" + + +@dataclass(frozen=True, slots=True) +class ParsedFetchUrl: + """Normalized URL pieces used for policy and the HTTP client.""" + + url: str + scheme: str + host: str + port: int + path_query: str # path + optional ?query (no fragment) + + +@dataclass(frozen=True, slots=True) +class HostAssessment: + host: str + port: int + classification: HostClass + addresses: tuple[str, ...] + reason: str + + +def normalize_method(method: str) -> str: + upper = method.strip().upper() + if upper not in ALLOWED_METHODS: + allowed = ", ".join(sorted(ALLOWED_METHODS)) + msg = f"method must be one of {allowed}; got {method!r}" + raise FetchPolicyError(msg) + return upper + + +def parse_fetch_url(url: str) -> ParsedFetchUrl: + raw = url.strip() + if not raw: + msg = "url must not be empty" + raise FetchPolicyError(msg) + parsed = urlparse(raw) + scheme = (parsed.scheme or "").lower() + if scheme not in ALLOWED_SCHEMES: + msg = f"url scheme must be http or https; got {scheme or '(none)'!r}" + raise FetchPolicyError(msg) + if not parsed.hostname: + msg = "url must include a hostname" + raise FetchPolicyError(msg) + host = parsed.hostname + # urlparse keeps brackets out of hostname for IPv6. + port = parsed.port + if port is None: + port = 443 if scheme == "https" else 80 + # Drop fragment; rebuild without params quirks. + path = parsed.path or "/" + query = parsed.query + path_query = path if not query else f"{path}?{query}" + # Prefer normalized form (no fragment). + normalized = urlunparse((scheme, parsed.netloc, path, "", query, "")) + return ParsedFetchUrl( + url=normalized, + scheme=scheme, + host=host, + port=port, + path_query=path_query, + ) + + +def resolve_redirect_url(current: str, location: str) -> str: + """Resolve a redirect Location against *current* and re-parse for safety.""" + if not location or not location.strip(): + msg = "redirect Location is empty" + raise FetchPolicyError(msg) + joined = urljoin(current, location.strip()) + return parse_fetch_url(joined).url + + +def _ip_classification(address: str) -> HostClass: + try: + ip = ipaddress.ip_address(address) + except ValueError: + return HostClass.FORBIDDEN + for network in _NEVER_ALLOW_NETWORKS: + if ip in network: + return HostClass.FORBIDDEN + # IPv6 unique-local / link-local / loopback / unspecified + if ip.is_loopback or ip.is_link_local or ip.is_private or ip.is_reserved or ip.is_multicast or ip.is_unspecified: + return HostClass.PRIVATE + return HostClass.PUBLIC + + +def classify_ip_strings(addresses: Sequence[str]) -> HostClass: + """Worst-class wins: FORBIDDEN > PRIVATE > PUBLIC.""" + if not addresses: + return HostClass.FORBIDDEN + worst = HostClass.PUBLIC + for addr in addresses: + kind = _ip_classification(addr) + if kind is HostClass.FORBIDDEN: + return HostClass.FORBIDDEN + if kind is HostClass.PRIVATE: + worst = HostClass.PRIVATE + return worst + + +async def resolve_host_addresses(host: str) -> tuple[str, ...]: + """Resolve *host* via the running loop's ``getaddrinfo`` (async).""" + # Literal IPs: no DNS. + try: + ip = ipaddress.ip_address(host) + except ValueError: + ip = None + if ip is not None: + return (str(ip),) + + import asyncio + + loop = asyncio.get_running_loop() + try: + infos = await loop.getaddrinfo( + host, + None, + type=socket.SOCK_STREAM, + proto=socket.IPPROTO_TCP, + ) + except socket.gaierror as exc: + msg = f"DNS resolution failed for {host!r}: {exc}" + raise FetchPolicyError(msg) from exc + addrs: list[str] = [] + seen: set[str] = set() + for info in infos: + sockaddr = info[4] + if not sockaddr: + continue + addr = str(sockaddr[0]) + if addr not in seen: + seen.add(addr) + addrs.append(addr) + if not addrs: + msg = f"DNS resolution returned no addresses for {host!r}" + raise FetchPolicyError(msg) + return tuple(addrs) + + +async def assess_host(host: str, port: int) -> HostAssessment: + """Classify *host* after resolution (literal IP or DNS).""" + addresses = await resolve_host_addresses(host) + classification = classify_ip_strings(addresses) + if classification is HostClass.FORBIDDEN: + reason = f"host {host!r} resolves to a forbidden address ({', '.join(addresses)})" + elif classification is HostClass.PRIVATE: + reason = f"host {host!r} is loopback/private/link-local ({', '.join(addresses)})" + else: + reason = f"host {host!r} is public ({', '.join(addresses)})" + return HostAssessment( + host=host, + port=port, + classification=classification, + addresses=addresses, + reason=reason, + ) + + +def grant_key(host: str, port: int) -> str: + """Stable key for instance-scoped private fetch grants.""" + return f"{host.lower()}:{port}" + + +def normalize_request_headers( + headers: Mapping[str, str] | None, + *, + user_agent: str | None = None, +) -> dict[str, str]: + """Build outbound headers. + + * Non-empty *user_agent* wins and replaces any User-Agent in *headers*. + * Else keep a User-Agent already present in *headers* (never replaced by default). + * When both omit UA, set :data:`DEFAULT_USER_AGENT`. + * Forbidden hop-by-hop headers raise :class:`FetchPolicyError`. + """ + out: dict[str, str] = {} + if headers: + for raw_key, raw_value in headers.items(): + key = raw_key.strip() + if not key: + msg = "header name must not be empty" + raise FetchPolicyError(msg) + lower = key.lower() + if lower in _FORBIDDEN_REQUEST_HEADERS: + msg = f"header {key!r} is not allowed" + raise FetchPolicyError(msg) + out[key] = raw_value + + def _has_user_agent(mapping: Mapping[str, str]) -> bool: + return any(k.lower() == "user-agent" for k in mapping) + + # Tool-call User-Agent is never replaced by a host default. Prefer the + # dedicated ``user_agent`` argument when non-empty; else keep headers; else default. + if user_agent is not None and user_agent.strip(): + out = {key: value for key, value in out.items() if key.lower() != "user-agent"} + out["User-Agent"] = user_agent.strip() + elif not _has_user_agent(out): + out["User-Agent"] = DEFAULT_USER_AGENT + return out + + +def is_https_to_http_downgrade(previous_scheme: str, next_scheme: str) -> bool: + return previous_scheme.lower() == "https" and next_scheme.lower() == "http" + + +def soft_confirm_reason( + *, + method: str, + url: str, + scheme: str, + body_present: bool, +) -> str | None: + """Soft-confirm (YOLO-eligible only when reason is None or caller tags allow). + + Public cleartext HTTP and any mutating method get a soft reason. + HTTPS GET/HEAD-like GET with no body: no soft reason (still subject to caps). + """ + parts: list[str] = [f"fetch: {method} {url}"] + risky = False + if scheme.lower() == "http": + parts.append("cleartext HTTP (not HTTPS)") + risky = True + if method in {"POST", "PUT", "DELETE"}: + parts.append(f"mutating method {method}") + if body_present: + parts.append("request body present") + risky = True + if not risky: + return None + return "\n ".join(parts)