Harden web UI security

1 files changed, 7 insertions, 59 deletions

diff --git a/src/ollama_web/tools/pdf.py b/src/ollama_web/tools/pdf.py
index 815c053..373506c 100644
--- a/src/ollama_web/tools/pdf.py
+++ b/src/ollama_web/tools/pdf.py
@@ -1,48 +1,10 @@
-"""PDF text extraction helpers."""
+"""PDF tool exposed to Ollama."""
 
 from __future__ import annotations
 
-import io
-from pathlib import Path
-
-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.
-
-    Args:
-      data: Raw PDF file bytes.
-      max_chars: Maximum characters to return (default from settings).
-
-    Returns:
-      Extracted plain text. Errors are returned as short strings.
-    """
-    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
-
-    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:
-        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."
+from .helper.pdf import extract_pdf_text
+from .helper.safe_http import UnsafeURL, fetch_public_bytes
 
 
 def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str:
@@ -61,28 +23,14 @@ def pdf_to_text(url_or_path: str, max_chars: int | None = None) -> str:
     src = url_or_path.strip()
     data: bytes | None = None
     if src.startswith("http://") or src.startswith("https://"):
-        import httpx
-        from ..config import settings as cfg  # local import to avoid top-level cycles
-
         try:
-            with httpx.Client(
-                timeout=cfg.scrape_timeout,
-                follow_redirects=True,
-                headers={"User-Agent": "ollama-web/0.1 (+https://github.com/local)"},
-            ) as client:
-                resp = client.get(src)
-                resp.raise_for_status()
-                data = resp.content
+            data, _, _ = fetch_public_bytes(src, timeout=settings.scrape_timeout)
+        except UnsafeURL as exc:
+            return f"Blocked URL: {exc}"
         except Exception as exc:  # noqa: BLE001
             return f"Failed to fetch PDF: {exc}"
     else:
-        path = Path(src)
-        if not path.exists():
-            return f"PDF not found: {src}"
-        try:
-            data = path.read_bytes()
-        except Exception as exc:  # noqa: BLE001
-            return f"Failed to read PDF: {exc}"
+        return "Only http(s) PDF URLs are supported."
 
     if data is None:
         return "No PDF data loaded."