atc-agents/api/auth.py

"""Authentik OIDC login and session cookie for Command Center."""

from __future__ import annotations

import logging
import os
from typing import Any
from urllib.parse import quote, urlencode

import httpx
from authlib.common.security import generate_token
from authlib.integrations.starlette_client import OAuth
from authlib.oauth2.rfc7636 import create_s256_code_challenge
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, RedirectResponse
from starlette.middleware.sessions import SessionMiddleware

log = logging.getLogger("atc-agents.auth")

AUTH_ENABLED = os.getenv("AUTH_ENABLED", "true").lower() in ("1", "true", "yes")
AUTHENTIK_ISSUER = os.getenv(
    "AUTHENTIK_ISSUER",
    "http://atc-mgt01.dell-atc.lan:9000/application/o/command-center/",
).rstrip("/") + "/"
AUTHENTIK_CLIENT_ID = os.getenv("AUTHENTIK_CLIENT_ID", "")
AUTHENTIK_CLIENT_SECRET = os.getenv("AUTHENTIK_CLIENT_SECRET", "")
AUTHENTIK_REDIRECT_URI = os.getenv(
    "AUTHENTIK_REDIRECT_URI",
    "http://10.0.21.33/auth/callback",
)
SESSION_SECRET = os.getenv("SESSION_SECRET", "dev-insecure-change-me")
SESSION_COOKIE = "cc_session"

oauth = OAuth()
_oauth_ready = False

PUBLIC_PREFIXES = (
    "/auth/",
    "/api/health",
    "/api/auth/me",
)


def is_auth_enabled() -> bool:
    return AUTH_ENABLED


def build_session_user(claims: dict[str, Any]) -> dict[str, Any]:
    name = claims.get("name") or claims.get("preferred_username") or claims.get("email") or "user"
    email = claims.get("email") or ""
    username = claims.get("preferred_username") or claims.get("nickname") or email or str(claims.get("sub") or "user")
    return {
        "sub": claims.get("sub"),
        "email": email,
        "name": name,
        "preferred_username": username,
    }


def get_session_user(request: Request) -> dict[str, Any] | None:
    if not AUTH_ENABLED:
        return {
            "sub": "dev:local",
            "email": "",
            "name": "Dev User",
            "preferred_username": "dev",
            "dev": True,
        }
    user = request.session.get("user")
    return user if isinstance(user, dict) else None


def auth_me_payload(request: Request) -> dict[str, Any]:
    user = get_session_user(request)
    if not user:
        raise HTTPException(status_code=401, detail="Not authenticated")
    return {
        "user": user.get("sub"),
        "email": user.get("email"),
        "name": user.get("name"),
        "preferred_username": user.get("preferred_username"),
        "auth_enabled": AUTH_ENABLED,
    }


def _ensure_oauth() -> None:
    global _oauth_ready
    if _oauth_ready or not AUTH_ENABLED:
        return
    if not AUTHENTIK_CLIENT_ID or not AUTHENTIK_CLIENT_SECRET:
        log.warning("AUTH_ENABLED but Authentik client credentials missing")
        return
    meta_url = AUTHENTIK_ISSUER + ".well-known/openid-configuration"
    oauth.register(
        name="authentik",
        client_id=AUTHENTIK_CLIENT_ID,
        client_secret=AUTHENTIK_CLIENT_SECRET,
        server_metadata_url=meta_url,
        client_kwargs={"scope": "openid profile email"},
    )
    _oauth_ready = True


def _is_public_path(path: str) -> bool:
    return any(path == p or path.startswith(p) for p in PUBLIC_PREFIXES)


async def auth_guard_middleware(request: Request, call_next):
    if not AUTH_ENABLED:
        return await call_next(request)
    path = request.url.path
    if _is_public_path(path):
        return await call_next(request)
    if path.startswith("/api/"):
        if get_session_user(request) is None:
            return JSONResponse({"detail": "Not authenticated"}, status_code=401)
    return await call_next(request)


def init_auth_middleware(app: FastAPI) -> None:
    """Session must wrap auth guard so request.session is populated first."""
    app.middleware("http")(auth_guard_middleware)
    app.add_middleware(
        SessionMiddleware,
        secret_key=SESSION_SECRET,
        session_cookie=SESSION_COOKIE,
        max_age=86400 * 7,
        same_site="lax",
        https_only=False,
    )


def setup_auth(app: FastAPI) -> None:
    _ensure_oauth()
    _register_routes(app)


def _token_endpoint() -> str:
    return AUTHENTIK_ISSUER.rstrip("/").rsplit("/application/o/", 1)[0] + "/application/o/token/"


def _userinfo_endpoint() -> str:
    return AUTHENTIK_ISSUER.rstrip("/").rsplit("/application/o/", 1)[0] + "/application/o/userinfo/"


async def _exchange_code_for_userinfo(code: str, code_verifier: str | None) -> dict[str, Any]:
    data = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": AUTHENTIK_REDIRECT_URI,
        "client_id": AUTHENTIK_CLIENT_ID,
        "client_secret": AUTHENTIK_CLIENT_SECRET,
    }
    if code_verifier:
        data["code_verifier"] = code_verifier
    async with httpx.AsyncClient(timeout=20.0) as client:
        tok = await client.post(_token_endpoint(), data=data)
        if tok.status_code >= 400:
            log.warning("Token exchange failed: %s %s", tok.status_code, tok.text[:300])
            tok.raise_for_status()
        payload = tok.json()
        access = payload.get("access_token")
        if not access:
            raise RuntimeError("No access_token in token response")
        ui = await client.get(
            _userinfo_endpoint(),
            headers={"Authorization": f"Bearer {access}"},
        )
        if ui.status_code >= 400:
            log.warning("Userinfo failed: %s %s", ui.status_code, ui.text[:300])
            ui.raise_for_status()
        return ui.json()


def _register_routes(app: FastAPI) -> None:
    @app.get("/api/auth/me")
    async def api_auth_me(request: Request):
        return auth_me_payload(request)

    @app.get("/auth/login")
    async def auth_login(request: Request):
        if not AUTH_ENABLED:
            return RedirectResponse("/", status_code=302)
        _ensure_oauth()
        if not _oauth_ready:
            raise HTTPException(503, "Authentik not configured")
        code_verifier = generate_token(48)
        request.session["pkce_code_verifier"] = code_verifier
        return await oauth.authentik.authorize_redirect(
            request,
            AUTHENTIK_REDIRECT_URI,
            code_challenge=create_s256_code_challenge(code_verifier),
            code_challenge_method="S256",
            code_verifier=code_verifier,
        )

    @app.get("/auth/callback")
    async def auth_callback(request: Request):
        if not AUTH_ENABLED:
            return RedirectResponse("/", status_code=302)
        err = request.query_params.get("error")
        if err:
            desc = request.query_params.get("error_description") or err
            log.warning("OIDC provider error: %s — %s", err, desc)
            return RedirectResponse(f"/?error={quote(desc)}", status_code=302)
        code = request.query_params.get("code")
        if not code:
            return RedirectResponse("/?error=missing_code", status_code=302)
        state = request.query_params.get("state")
        if state:
            request.session.pop(f"_state_authentik_{state}", None)
        code_verifier = request.session.pop("pkce_code_verifier", None)
        try:
            userinfo = await _exchange_code_for_userinfo(code, code_verifier)
        except Exception as e:
            log.warning("OIDC callback failed: %s", e)
            return RedirectResponse("/?error=login_failed", status_code=302)
        request.session["user"] = build_session_user(userinfo or {})
        return RedirectResponse("/", status_code=302)

    @app.get("/auth/logout")
    async def auth_logout(request: Request):
        request.session.clear()
        if not AUTH_ENABLED:
            return RedirectResponse("/", status_code=302)
        end_session = AUTHENTIK_ISSUER + "end-session/"
        post_logout = "http://10.0.21.33/"
        params = urlencode({"post_logout_redirect_uri": post_logout})
        return RedirectResponse(f"{end_session}?{params}", status_code=302)