open-swe/agent/utils/multimodal.py
Adam Moussa 2028208237
fix(security): SSRF-harden server-side image fetch + stop logging raw image URLs
INJ-01 (high): fetch_image_block used follow_redirects=True with no per-hop
revalidation and discarded the resolved-IP pin, so an attacker-authored Slack/
Linear image URL could 302-redirect the fetch to an internal host / cloud
metadata endpoint (blind SSRF), and DNS-rebinding could bypass the one-shot
is_url_safe check. Route image fetches through the same per-hop resolve+pin+
revalidate loop the http_request tool uses, lifted into url_safety as the shared
request_with_safe_redirects. Also strip the per-host Slack/Linear bearer token
on redirect so it can't be replayed to a redirect target.

SC-1 (low): linear.py logged full image URLs (which can carry signed tokens) at
DEBUG; multimodal logged them at INFO on every fetch. Log host-only.

Sink lived in multimodal.py (unchanged by the feature work) but PR #128 widened
its reach by no longer dropping images for text-only models. Fixing on the base
branch so #130/#129 inherit it on rebase. Adds fetch_image_block SSRF regression
tests (redirect-to-internal blocked; auth stripped on redirect).
2026-07-08 18:18:07 -04:00

116 lines
4.2 KiB
Python

"""Utilities for building multimodal content blocks."""
from __future__ import annotations
import base64
import logging
import mimetypes
import os
import re
from typing import Any
from urllib.parse import urlparse
import httpx
from langchain_core.messages.content import create_image_block
from .url_safety import request_with_safe_redirects
logger = logging.getLogger(__name__)
IMAGE_MARKDOWN_RE = re.compile(r"!\[[^\]]*\]\((https?://[^\s)]+)\)")
IMAGE_URL_RE = re.compile(
r"(https?://[^\s)]+\.(?:png|jpe?g|gif|webp|bmp|tiff)(?:\?[^\s)]+)?)",
re.IGNORECASE,
)
def extract_image_urls(text: str) -> list[str]:
"""Extract image URLs from markdown image syntax and direct image links."""
if not text:
return []
urls: list[str] = []
urls.extend(IMAGE_MARKDOWN_RE.findall(text))
urls.extend(IMAGE_URL_RE.findall(text))
deduped = dedupe_urls(urls)
if deduped:
logger.debug("Extracted %d image URL(s)", len(deduped))
return deduped
def vision_not_supported_warning(model_id: str, image_count: int) -> str:
"""Build a prompt-visible warning when images are sent to a text-only model."""
return (
f"\n\n**Note:** {image_count} image(s) were attached but the current model "
f"({model_id}) does not support image input. The images were not included. "
"Please switch to a vision-enabled model to process images."
)
async def fetch_image_block(
image_url: str,
client: httpx.AsyncClient,
) -> dict[str, Any] | None:
"""Fetch image bytes and build an image content block.
The fetch validates and pins every redirect hop (SSRF guard) and drops any
Authorization header once the URL redirects, so a per-host token is never
replayed to a redirect target the caller never chose to authenticate to.
URLs are logged host-only — a signed image URL can carry a bearer token.
"""
host = (urlparse(image_url).hostname or "").lower()
try:
headers = None
if host == "uploads.linear.app" or host.endswith(".uploads.linear.app"):
linear_api_key = os.environ.get("LINEAR_API_KEY", "")
if linear_api_key:
headers = {"Authorization": linear_api_key}
else:
logger.warning(
"LINEAR_API_KEY not set; cannot authenticate image fetch for %s", host
)
elif host == "files.slack.com" or host.endswith(".files.slack.com"):
slack_bot_token = os.environ.get("SLACK_BOT_TOKEN", "")
if slack_bot_token:
headers = {"Authorization": f"Bearer {slack_bot_token}"}
else:
logger.warning(
"SLACK_BOT_TOKEN not set; cannot authenticate image fetch for %s", host
)
response, blocked = await request_with_safe_redirects(
client, "GET", image_url, headers=headers, strip_auth_on_redirect=True
)
if blocked is not None:
_, reason = blocked
logger.warning("Refusing to fetch image (SSRF guard) from %s: %s", host, reason)
return None
response.raise_for_status()
content_type = response.headers.get("Content-Type", "").split(";")[0].strip()
if not content_type:
guessed, _ = mimetypes.guess_type(image_url)
if not guessed:
logger.warning("Could not determine content type from %s; skipping image", host)
return None
content_type = guessed
supported_types = {"image/jpeg", "image/png", "image/gif", "image/webp"}
if content_type not in supported_types:
logger.warning(
"Unsupported content type '%s' from %s; skipping image", content_type, host
)
return None
encoded = base64.b64encode(response.content).decode("ascii")
logger.info(
"Fetched image from %s (%s, %d bytes)", host, content_type, len(response.content)
)
return create_image_block(base64=encoded, mime_type=content_type)
except Exception:
logger.exception("Failed to fetch image from %s", host)
return None
def dedupe_urls(urls: list[str]) -> list[str]:
return list(dict.fromkeys(urls))