Initial Commit

app/owuitools/security.py

@@ -0,0 +1,155 @@
+"""Request authentication and URL sanity checks.
+
+Read the docstring in configs/services/egress.squid.conf before changing
+anything here.
+
+The checks in this module are a *usability* layer, not the security
+boundary. They exist to turn "the model asked for something impossible"
+into a fast, readable error the model can correct from, instead of a
+30-second timeout. The actual boundary is that this process has no route
+to anything except squid, and squid denies private destinations after
+resolving them itself.
+
+That distinction matters: because every request is proxied, this process
+never opens the socket, so it cannot pin a connection to a
+previously-validated IP. Anything that tried to enforce security here
+would be enforcing it in the wrong place.
+"""
+
+from __future__ import annotations
+
+import ipaddress
+import re
+from urllib.parse import urlparse, urlsplit, urlunsplit
+
+from fastapi import Header, HTTPException, status
+
+from .settings import settings
+
+ALLOWED_SCHEMES = {"http", "https"}
+ALLOWED_PORTS = {80, 443, None}
+
+# Names that never point anywhere legitimate from a container with no
+# route to the local network. Blocked here purely so the failure is
+# instant and self-explanatory.
+_BLOCKED_SUFFIXES = (
+    ".local", ".localdomain", ".internal", ".intranet", ".lan",
+    ".home", ".corp", ".private", ".onion",
+)
+_BLOCKED_NAMES = {
+    "localhost", "metadata", "metadata.google.internal",
+    "instance-data", "host.docker.internal",
+}
+
+
+class UnsafeURL(ValueError):
+    """Raised with a message written to be read by the model."""
+
+
+def require_auth(authorization: str | None = Header(default=None, include_in_schema=False)) -> None:
+    """Bearer check for every tool endpoint.
+
+    Open WebUI sends `Authorization: Bearer <key>` from the connection's
+    stored key. /health is deliberately left open so Docker can probe it.
+    """
+    expected = f"Bearer {settings.auth_token}"
+    if not authorization or authorization != expected:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or missing bearer token for the owuitools server.",
+        )
+
+
+def validate_url(raw: str, *, allow_http: bool = True) -> str:
+    """Normalise a model-supplied URL, or explain why it cannot be used."""
+    if not raw or not raw.strip():
+        raise UnsafeURL("No URL was provided.")
+
+    url = raw.strip()
+    # Models routinely hand over bare hostnames.
+    if "://" not in url:
+        url = f"https://{url}"
+
+    try:
+        parts = urlsplit(url)
+    except ValueError as exc:
+        raise UnsafeURL(f"That is not a parseable URL: {exc}") from None
+
+    scheme = parts.scheme.lower()
+    if scheme not in ALLOWED_SCHEMES:
+        raise UnsafeURL(
+            f"Only http and https URLs can be fetched, not '{scheme}:'. "
+            "There is no way to read local files, FTP, or gopher with this tool."
+        )
+    if scheme == "http" and not allow_http:
+        raise UnsafeURL("This tool requires an https URL.")
+
+    host = (parts.hostname or "").lower().rstrip(".")
+    if not host:
+        raise UnsafeURL(f"No hostname found in '{raw}'.")
+
+    try:
+        port = parts.port
+    except ValueError:
+        raise UnsafeURL(f"Invalid port in '{raw}'.") from None
+    if port not in ALLOWED_PORTS:
+        raise UnsafeURL(
+            f"Port {port} cannot be reached. Only standard web ports 80 and 443 are available."
+        )
+
+    _reject_private_literal(host)
+    _reject_internal_name(host)
+
+    # Strip fragments; they are never sent and only confuse cache keys.
+    return urlunsplit((scheme, parts.netloc, parts.path or "/", parts.query, ""))
+
+
+def _reject_private_literal(host: str) -> None:
+    literal = host[1:-1] if host.startswith("[") and host.endswith("]") else host
+    try:
+        ip = ipaddress.ip_address(literal)
+    except ValueError:
+        return  # not an IP literal; squid will resolve and judge the name
+
+    if not ip.is_global or ip.is_multicast:
+        raise UnsafeURL(
+            f"{ip} is a private, loopback, link-local or otherwise non-public address. "
+            "This tool can only reach the public internet. If the user asked about a "
+            "device on their own network, say that you have no access to it rather than "
+            "guessing what it might return."
+        )
+
+
+def _reject_internal_name(host: str) -> None:
+    if host in _BLOCKED_NAMES or host.endswith(_BLOCKED_SUFFIXES):
+        raise UnsafeURL(
+            f"'{host}' is a private/internal hostname that this tool cannot reach. "
+            "Only public internet addresses are available."
+        )
+    # A single-label name ("router", "nas") can only resolve inside a local
+    # network, so it is never a valid target from here.
+    if "." not in host:
+        raise UnsafeURL(
+            f"'{host}' is not a public domain name. Provide a full domain such as example.com."
+        )
+
+
+_HOSTNAME_RE = re.compile(r"^(?=.{1,253}$)(?!-)[A-Za-z0-9-]{1,63}(?<!-)(\.(?!-)[A-Za-z0-9-]{1,63}(?<!-))+$")
+
+
+def validate_hostname(raw: str) -> str:
+    """Validate a bare hostname for the probe tools (DNS, TLS, RDAP)."""
+    if not raw or not raw.strip():
+        raise UnsafeURL("No hostname was provided.")
+
+    host = raw.strip().lower().rstrip(".")
+    # Tolerate a pasted URL where a hostname was asked for.
+    if "://" in host:
+        host = (urlparse(host).hostname or "").lower().rstrip(".")
+
+    _reject_private_literal(host)
+    _reject_internal_name(host)
+
+    if not _HOSTNAME_RE.match(host):
+        raise UnsafeURL(f"'{raw}' is not a valid public hostname.")
+    return host