Harden web UI security

4 files changed, 221 insertions, 0 deletions

diff --git a/src/ollama_web/tools/helper/__init__.py b/src/ollama_web/tools/helper/__init__.py
new file mode 100644
index 0000000..f201101
--- /dev/null
+++ b/src/ollama_web/tools/helper/__init__.py
@@ -0,0 +1,13 @@
+"""Internal helper functions for tools."""
+
+from .image import resize_image
+from .pdf import extract_pdf_text
+from .safe_http import UnsafeURL, fetch_public_bytes, validate_public_http_url
+
+__all__ = [
+    "UnsafeURL",
+    "extract_pdf_text",
+    "fetch_public_bytes",
+    "resize_image",
+    "validate_public_http_url",
+]
diff --git a/src/ollama_web/tools/helper/image.py b/src/ollama_web/tools/helper/image.py
new file mode 100644
index 0000000..e89a96f
--- /dev/null
+++ b/src/ollama_web/tools/helper/image.py
@@ -0,0 +1,65 @@
+"""Image preprocessing helpers for ollama vision models."""
+
+from __future__ import annotations
+
+import io
+from pathlib import Path
+
+from PIL import Image
+
+from ...config import settings
+
+Image.MAX_IMAGE_PIXELS = settings.max_image_pixels
+
+
+_IMAGE_FORMATS: dict[str, str] = {
+    "png": "PNG",
+    "jpg": "JPEG",
+    "jpeg": "JPEG",
+    "gif": "GIF",
+    "bmp": "BMP",
+    "webp": "WEBP",
+    "tif": "TIFF",
+    "tiff": "TIFF",
+}
+
+
+def _format_from_name(name: str) -> str:
+    ext = Path(name).suffix.lstrip(".").lower()
+    return _IMAGE_FORMATS.get(ext, "JPEG")
+
+
+def resize_image(
+    data: bytes,
+    name: str = "image.jpg",
+    max_dimension: int = 1024,
+    quality: int = 85,
+) -> bytes:
+    """Resize and re-encode an image to reduce payload size for ollama."""
+    if not data:
+        return data
+
+    try:
+        with Image.open(io.BytesIO(data)) as img:
+            if img.width * img.height > settings.max_image_pixels:
+                return data
+            processed = img.convert("RGB")
+
+            width, height = processed.size
+            if width > max_dimension or height > max_dimension:
+                ratio = min(max_dimension / width, max_dimension / height)
+                new_size = (int(width * ratio), int(height * ratio))
+                processed = processed.resize(new_size, Image.Resampling.LANCZOS)
+
+            fmt = _format_from_name(name)
+            output = io.BytesIO()
+            save_kwargs: dict[str, object] = {}
+            if fmt == "JPEG":
+                save_kwargs["quality"] = quality
+                save_kwargs["optimize"] = True
+            elif fmt == "WEBP":
+                save_kwargs["quality"] = quality
+            processed.save(output, format=fmt, **save_kwargs)
+            return output.getvalue()
+    except Exception:  # noqa: BLE001
+        return data
diff --git a/src/ollama_web/tools/helper/pdf.py b/src/ollama_web/tools/helper/pdf.py
new file mode 100644
index 0000000..51d3b3a
--- /dev/null
+++ b/src/ollama_web/tools/helper/pdf.py
@@ -0,0 +1,39 @@
+"""PDF text extraction helpers."""
+
+from __future__ import annotations
+
+import io
+
+from pypdf import PdfReader
+
+from ...config import settings
+
+
+def extract_pdf_text(data: bytes, max_chars: int | None = None) -> str:
+    """Extract text from PDF bytes."""
+    if not data:
+        return "Empty PDF data."
+
+    limit = max_chars if max_chars is not None and max_chars > 0 else settings.fetch_max_chars
+
+    if len(data) > settings.max_fetch_mb * 1024 * 1024:
+        return "PDF exceeds the configured size limit."
+
+    try:
+        reader = PdfReader(stream=io.BytesIO(data))
+    except Exception as exc:  # noqa: BLE001
+        return f"Failed to parse PDF: {exc}"
+
+    parts: list[str] = []
+    for page in reader.pages[: settings.max_pdf_pages]:
+        try:
+            text = page.extract_text()
+            if text:
+                parts.append(text)
+        except Exception:  # noqa: BLE001
+            continue
+
+    full = "\n\n".join(parts)
+    if len(full) > limit:
+        full = full[:limit].rsplit(" ", 1)[0] + "\n…[truncated]"
+    return full or "No text could be extracted from the PDF."
diff --git a/src/ollama_web/tools/helper/safe_http.py b/src/ollama_web/tools/helper/safe_http.py
new file mode 100644
index 0000000..17157db
--- /dev/null
+++ b/src/ollama_web/tools/helper/safe_http.py
@@ -0,0 +1,104 @@
+"""Safe outbound HTTP fetching for LLM tools."""
+
+from __future__ import annotations
+
+import ipaddress
+import socket
+
+import httpx
+
+from ...config import settings
+
+_METADATA_IPS = {
+    ipaddress.ip_address("169.254.169.254"),
+}
+
+
+class UnsafeURL(ValueError):
+    """Raised when an outbound URL is not allowed."""
+
+
+def _is_public_ip(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
+    if ip in _METADATA_IPS:
+        return False
+    return not (
+        ip.is_private
+        or ip.is_loopback
+        or ip.is_link_local
+        or ip.is_multicast
+        or ip.is_reserved
+        or ip.is_unspecified
+    )
+
+
+def validate_public_http_url(url: str) -> httpx.URL:
+    """Return a parsed URL after rejecting private network targets."""
+    parsed = httpx.URL(url)
+    if parsed.scheme not in {"http", "https"}:
+        raise UnsafeURL("Only http(s) URLs are allowed.")
+    host = parsed.host
+    if not host:
+        raise UnsafeURL("URL host is required.")
+    if host.lower() == "localhost" or host.lower().endswith(".localhost"):
+        raise UnsafeURL("Localhost URLs are not allowed.")
+
+    try:
+        ip = ipaddress.ip_address(host)
+        if not _is_public_ip(ip):
+            raise UnsafeURL("Private or local IP addresses are not allowed.")
+        return parsed
+    except ValueError:
+        pass
+
+    try:
+        infos = socket.getaddrinfo(
+            host,
+            parsed.port or (443 if parsed.scheme == "https" else 80),
+            type=socket.SOCK_STREAM,
+        )
+    except socket.gaierror as exc:
+        raise UnsafeURL(f"Could not resolve host: {host}") from exc
+
+    addresses = {item[4][0] for item in infos}
+    if not addresses:
+        raise UnsafeURL(f"Could not resolve host: {host}")
+    for address in addresses:
+        ip = ipaddress.ip_address(address)
+        if not _is_public_ip(ip):
+            raise UnsafeURL("Resolved private or local IP addresses are not allowed.")
+    return parsed
+
+
+def fetch_public_bytes(
+    url: str,
+    *,
+    max_bytes: int | None = None,
+    timeout: float | None = None,
+) -> tuple[bytes, str, str]:
+    """Fetch bytes from a public http(s) URL with redirect and size checks."""
+    limit = max_bytes or settings.max_fetch_mb * 1024 * 1024
+    current = validate_public_http_url(url)
+    headers = {"User-Agent": "ollama-web/0.1 (+https://github.com/local)"}
+    with httpx.Client(timeout=timeout or settings.scrape_timeout, headers=headers) as client:
+        for _ in range(6):
+            with client.stream("GET", current, follow_redirects=False) as resp:
+                if resp.is_redirect:
+                    location = resp.headers.get("location")
+                    if not location:
+                        raise UnsafeURL("Redirect without Location header.")
+                    current = validate_public_http_url(str(current.join(location)))
+                    continue
+                resp.raise_for_status()
+                declared = resp.headers.get("content-length")
+                if declared and int(declared) > limit:
+                    raise UnsafeURL("Remote response exceeds the configured size limit.")
+                chunks: list[bytes] = []
+                total = 0
+                for chunk in resp.iter_bytes():
+                    total += len(chunk)
+                    if total > limit:
+                        raise UnsafeURL("Remote response exceeds the configured size limit.")
+                    chunks.append(chunk)
+                content_type = resp.headers.get("content-type", "")
+                return b"".join(chunks), str(resp.url), content_type
+        raise UnsafeURL("Too many redirects.")