FabledSteward/steward/auth/middleware.py

from __future__ import annotations
import functools
import time
from urllib.parse import quote, urlsplit
from quart import session, redirect, url_for, abort, g, request, Response
from steward.models.users import UserRole

_ROLE_ORDER = [UserRole.viewer, UserRole.operator, UserRole.admin]

# Paths we never bounce back to after login (would loop or make no sense).
_NO_RETURN = {"/login", "/logout", "/setup"}


def role_meets(user_role: UserRole, minimum_role: UserRole) -> bool:
    """True when user_role is at least minimum_role in the viewer<operator<admin order."""
    return _ROLE_ORDER.index(user_role) >= _ROLE_ORDER.index(minimum_role)


def safe_next_url(raw: str | None) -> str | None:
    """Return a safe same-site relative path+query from `raw`, else None.

    Accepts a relative path or a full URL (e.g. HTMX's HX-Current-URL) and
    reduces it to path[?query]. Rejects off-site/protocol-relative targets and
    the auth pages themselves, so it can't be abused as an open redirect.
    """
    if not raw:
        return None
    parts = urlsplit(raw)
    # Protocol-relative (//host/…) has a netloc but no scheme — reject it; a full
    # same-site URL (HX-Current-URL) has a scheme and we keep only its path.
    if parts.netloc and not parts.scheme:
        return None
    path = parts.path
    if parts.query:
        path += "?" + parts.query
    if not path.startswith("/") or path.startswith("//") or path.startswith("/\\"):
        return None
    if path.split("?", 1)[0] in _NO_RETURN:
        return None
    return path


def _login_redirect():
    """Redirect an unauthenticated request to /login, preserving where to return.

    For HTMX requests the destination is the page the user is on (HX-Current-URL)
    and we use HX-Redirect so the whole browser navigates (not a fragment swap);
    for normal requests it's the requested path.
    """
    is_htmx = request.headers.get("HX-Request") == "true"
    nxt = safe_next_url(request.headers.get("HX-Current-URL")) if is_htmx else None
    if nxt is None:
        nxt = safe_next_url(request.full_path)
    login_url = url_for("auth.login")
    if nxt:
        login_url = f"{login_url}?next={quote(nxt, safe='/')}"
    if is_htmx:
        resp = Response("", status=401)
        resp.headers["HX-Redirect"] = login_url
        return resp
    return redirect(login_url)


def require_role(minimum_role: UserRole):
    """Decorator: requires authenticated user with at least minimum_role.
    Also allows access for validated share-token requests (viewer level only).
    """
    def decorator(f):
        @functools.wraps(f)
        async def wrapper(*args, **kwargs):
            # Share-token requests bypass session auth for viewer-level endpoints
            if getattr(g, "share_viewer", False) and minimum_role == UserRole.viewer:
                return await f(*args, **kwargs)
            user_id = session.get("user_id")
            if not user_id:
                return _login_redirect()
            user_role = UserRole(session.get("user_role", "viewer"))
            if _ROLE_ORDER.index(user_role) < _ROLE_ORDER.index(minimum_role):
                abort(403)
            return await f(*args, **kwargs)
        wrapper.steward_role = minimum_role
        return wrapper
    return decorator


def public_route(reason: str):
    """Mark a view that is deliberately reachable without a role — login,
    the health probe, a bearer-authenticated endpoint. Every other view must
    carry require_role; check_view_is_marked refuses one that has neither, so
    a route that forgets its decorator is closed rather than open."""
    def decorator(f):
        f.steward_public = reason
        return f
    return decorator


async def check_view_is_marked() -> Response | None:
    """before_request hook: default-deny for views with neither require_role
    nor public_route. Static files and unmatched paths (404) pass through."""
    from quart import current_app
    endpoint = request.endpoint
    if endpoint is None or endpoint == "static" or endpoint.endswith(".static"):
        return None
    view = current_app.view_functions.get(endpoint)
    if view is None or hasattr(view, "steward_role") or hasattr(view, "steward_public"):
        return None
    current_app.logger.error("Refused %s: view %r has no require_role or public_route",
                             request.path, endpoint)
    abort(403)


def login_user(user) -> None:
    """Write user identity into the session, with a fresh CSRF token."""
    from steward.auth.csrf import rotate
    rotate()
    session.permanent = True
    session["user_id"] = user.id
    session["user_role"] = user.role.value
    session["username"] = user.username
    session["login_at"] = time.time()


async def refresh_session_user() -> None:
    """before_request: a session speaks for the user as they are now.

    A login ends once Settings → General's session lifetime has passed since
    it began (not since the last request: the cookie is re-signed on every
    response, so the cookie's own age would let an active session live
    forever). A user who has been deleted or deactivated is logged out on
    their next request, and a role change applies at once.
    """
    user_id = session.get("user_id")
    if not user_id:
        return
    from quart import current_app
    from steward.models.users import User
    lifetime = current_app.config.get("PERMANENT_SESSION_LIFETIME")
    started = session.get("login_at") or 0
    if lifetime is not None and time.time() - started > lifetime.total_seconds():
        session.clear()
        return
    async with current_app.db_sessionmaker() as db:
        user = await db.get(User, user_id)
    if user is None or not user.is_active:
        session.clear()
        return
    if session.get("user_role") != user.role.value:
        session["user_role"] = user.role.value


def logout_user() -> None:
    session.clear()


def current_user_id() -> str | None:
    return session.get("user_id")