1 files changed, 142 insertions, 0 deletions
diff --git a/src/ollama_web/auth.py b/src/ollama_web/auth.py
new file mode 100644
index 0000000..ba2066b
--- /dev/null
+++ b/src/ollama_web/auth.py
@@ -0,0 +1,142 @@
+"""Cookie authentication and CSRF protection."""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import secrets
+import time
+from collections.abc import Awaitable, Callable
+from http import HTTPStatus
+
+from starlette.datastructures import URL
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import JSONResponse, RedirectResponse, Response
+
+SESSION_COOKIE = "ollama_web_session"
+CSRF_COOKIE = "ollama_web_csrf"
+AUTH_MAX_AGE_SECONDS = 60 * 60 * 24 * 14
+SAFE_METHODS = {"GET", "HEAD", "OPTIONS"}
+AUTH_EXEMPT_PREFIXES = ("/static/",)
+AUTH_EXEMPT_PATHS = {
+    "/login",
+    "/api/auth/login",
+    "/api/auth/status",
+}
+
+
+def _b64(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
+
+
+def _sign(secret_key: str, value: str) -> str:
+    sig = hmac.new(secret_key.encode(), value.encode(), hashlib.sha256).digest()
+    return _b64(sig)
+
+
+def make_session_cookie(secret_key: str) -> str:
+    issued_at = str(int(time.time()))
+    nonce = secrets.token_urlsafe(16)
+    payload = f"{issued_at}.{nonce}"
+    return f"{payload}.{_sign(secret_key, payload)}"
+
+
+def verify_session_cookie(secret_key: str, value: str | None) -> bool:
+    if not value:
+        return False
+    parts = value.split(".")
+    if len(parts) != 3:
+        return False
+    payload = ".".join(parts[:2])
+    if not hmac.compare_digest(parts[2], _sign(secret_key, payload)):
+        return False
+    try:
+        issued_at = int(parts[0])
+    except ValueError:
+        return False
+    return 0 <= time.time() - issued_at <= AUTH_MAX_AGE_SECONDS
+
+
+def set_auth_cookies(response: Response, secret_key: str) -> str:
+    csrf = secrets.token_urlsafe(24)
+    response.set_cookie(
+        SESSION_COOKIE,
+        make_session_cookie(secret_key),
+        httponly=True,
+        samesite="lax",
+        max_age=AUTH_MAX_AGE_SECONDS,
+    )
+    response.set_cookie(
+        CSRF_COOKIE,
+        csrf,
+        httponly=False,
+        samesite="lax",
+        max_age=AUTH_MAX_AGE_SECONDS,
+    )
+    return csrf
+
+
+def clear_auth_cookies(response: Response) -> None:
+    response.delete_cookie(SESSION_COOKIE)
+    response.delete_cookie(CSRF_COOKIE)
+
+
+def is_authenticated(request: Request) -> bool:
+    return verify_session_cookie(
+        request.app.state.settings.secret_key,
+        request.cookies.get(SESSION_COOKIE),
+    )
+
+
+def _is_exempt(path: str) -> bool:
+    return path in AUTH_EXEMPT_PATHS or any(
+        path.startswith(prefix) for prefix in AUTH_EXEMPT_PREFIXES
+    )
+
+
+def _wants_json(request: Request) -> bool:
+    return request.url.path.startswith("/api/")
+
+
+def _login_redirect(request: Request) -> RedirectResponse:
+    url = URL("/login").include_query_params(next=request.url.path)
+    return RedirectResponse(str(url), status_code=303)
+
+
+class AuthMiddleware(BaseHTTPMiddleware):
+    """Protect app routes with a signed cookie and CSRF token."""
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        path = request.url.path
+        if not _is_exempt(path) and not is_authenticated(request):
+            if _wants_json(request):
+                return JSONResponse(
+                    {"error": "Authentication required"},
+                    status_code=HTTPStatus.UNAUTHORIZED,
+                )
+            return _login_redirect(request)
+
+        if (
+            path.startswith("/api/")
+            and path not in AUTH_EXEMPT_PATHS
+            and request.method not in SAFE_METHODS
+        ):
+            token_cookie = request.cookies.get(CSRF_COOKIE)
+            token_header = request.headers.get("x-csrf-token")
+            if (
+                not token_cookie
+                or not token_header
+                or not hmac.compare_digest(token_cookie, token_header)
+            ):
+                return JSONResponse(
+                    {"error": "Invalid CSRF token"},
+                    status_code=HTTPStatus.FORBIDDEN,
+                )
+
+        return await call_next(request)