owui-tools/app/owuitools/security.py

"""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