cx-ui/backend/app/auth.py

"""Authentication: signed session cookies, local accounts, and OIDC.

Local accounts exist so the tool runs on a laptop and in a cluster that has not
been wired to SSO yet. When CX_OIDC_ENABLED is set, Authentik becomes the source
of truth: users are created on first login, and admin rights follow a group
claim rather than being set here.

The session itself is the same either way - a signed, expiring cookie - so the
rest of the app never has to care which provider a user came from.
"""
from __future__ import annotations

import base64
import datetime as dt
import hashlib
import hmac
import json
import secrets
import time
from typing import Any, Optional

import httpx
from fastapi import Depends, HTTPException, Request, status
from sqlalchemy.orm import Session

from .config import get_settings
from .db import get_db
from .models import AuthProvider, User

SESSION_COOKIE = "cx_session"
settings = get_settings()


# --- password hashing -------------------------------------------------------
# PBKDF2 from the standard library: no native build step in the image, and
# strong enough for a small internal user table.

def hash_password(password: str) -> str:
    salt = secrets.token_bytes(16)
    digest = hashlib.pbkdf2_hmac("sha256", password.encode(), salt, 240_000)
    return f"pbkdf2_sha256$240000${base64.b64encode(salt).decode()}${base64.b64encode(digest).decode()}"


def verify_password(password: str, stored: str) -> bool:
    try:
        algo, rounds, salt_b64, digest_b64 = stored.split("$")
        if algo != "pbkdf2_sha256":
            return False
        expected = base64.b64decode(digest_b64)
        actual = hashlib.pbkdf2_hmac("sha256", password.encode(),
                                     base64.b64decode(salt_b64), int(rounds))
        return hmac.compare_digest(expected, actual)
    except (ValueError, TypeError):
        return False


# --- session cookie ---------------------------------------------------------

def _sign(payload: bytes) -> str:
    mac = hmac.new(settings.secret_key.encode(), payload, hashlib.sha256).digest()
    return f"{base64.urlsafe_b64encode(payload).decode()}.{base64.urlsafe_b64encode(mac).decode()}"


def issue_session(user_id: int) -> str:
    payload = json.dumps({
        "uid": user_id,
        "exp": int(time.time()) + settings.session_hours * 3600,
    }).encode()
    return _sign(payload)


def read_session(token: str) -> Optional[int]:
    try:
        body_b64, mac_b64 = token.split(".")
        payload = base64.urlsafe_b64decode(body_b64)
        expected = hmac.new(settings.secret_key.encode(), payload, hashlib.sha256).digest()
        if not hmac.compare_digest(expected, base64.urlsafe_b64decode(mac_b64)):
            return None
        data = json.loads(payload)
        if int(data.get("exp", 0)) < time.time():
            return None
        return int(data["uid"])
    except (ValueError, TypeError, KeyError, json.JSONDecodeError):
        return None


# --- dependencies -----------------------------------------------------------

def current_user(request: Request, db: Session = Depends(get_db)) -> User:
    token = request.cookies.get(SESSION_COOKIE, "")
    uid = read_session(token) if token else None
    if not uid:
        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Not signed in")
    user = db.get(User, uid)
    if not user or not user.is_active:
        raise HTTPException(status.HTTP_401_UNAUTHORIZED, "Account is not active")
    return user


def require_admin(user: User = Depends(current_user)) -> User:
    if not user.is_admin:
        raise HTTPException(status.HTTP_403_FORBIDDEN, "Administrator access required")
    return user


def optional_user(request: Request, db: Session = Depends(get_db)) -> Optional[User]:
    token = request.cookies.get(SESSION_COOKIE, "")
    uid = read_session(token) if token else None
    return db.get(User, uid) if uid else None


# --- local accounts ---------------------------------------------------------

def authenticate_local(db: Session, email: str, password: str) -> Optional[User]:
    if not settings.auth_local_enabled:
        return None
    user = db.query(User).filter(User.email == email.strip().lower()).one_or_none()
    if not user or not user.is_active or user.provider != AuthProvider.LOCAL:
        return None
    if not user.password_hash or not verify_password(password, user.password_hash):
        return None
    return user


def ensure_bootstrap_admin(db: Session) -> Optional[str]:
    """Create the first admin so a fresh deployment is reachable.

    Only ever runs when the user table is empty, and only with a password
    supplied through the environment - it never invents one.
    """
    if db.query(User).count():
        return None
    if not settings.bootstrap_admin_password:
        return ("No users exist and CX_BOOTSTRAP_ADMIN_PASSWORD is unset. "
                "Set it and restart, or sign in through SSO.")
    db.add(User(
        email=settings.bootstrap_admin_email.strip().lower(),
        name="Bootstrap admin",
        is_admin=True,
        provider=AuthProvider.LOCAL,
        password_hash=hash_password(settings.bootstrap_admin_password),
    ))
    db.commit()
    return f"Created bootstrap admin {settings.bootstrap_admin_email}"


# --- OIDC / Authentik -------------------------------------------------------

class OIDCError(RuntimeError):
    pass


_discovery_cache: dict[str, Any] = {}


async def oidc_discovery() -> dict[str, Any]:
    """Fetch and cache the provider metadata."""
    if not settings.oidc_issuer:
        raise OIDCError("CX_OIDC_ISSUER is not set.")
    if _discovery_cache.get("_issuer") == settings.oidc_issuer:
        return _discovery_cache
    url = settings.oidc_issuer.rstrip("/") + "/.well-known/openid-configuration"
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.get(url)
    if resp.status_code != 200:
        raise OIDCError(f"OIDC discovery failed ({resp.status_code}) at {url}")
    data = resp.json()
    data["_issuer"] = settings.oidc_issuer
    _discovery_cache.clear()
    _discovery_cache.update(data)
    return data


def oidc_state() -> str:
    """A signed, short-lived value tying the callback to this browser."""
    payload = json.dumps({"n": secrets.token_urlsafe(16), "exp": int(time.time()) + 600}).encode()
    return _sign(payload)


def oidc_state_valid(state: str) -> bool:
    try:
        body_b64, mac_b64 = state.split(".")
        payload = base64.urlsafe_b64decode(body_b64)
        expected = hmac.new(settings.secret_key.encode(), payload, hashlib.sha256).digest()
        if not hmac.compare_digest(expected, base64.urlsafe_b64decode(mac_b64)):
            return False
        return int(json.loads(payload).get("exp", 0)) >= time.time()
    except (ValueError, TypeError, KeyError, json.JSONDecodeError):
        return False


async def oidc_exchange(code: str, redirect_uri: str) -> dict[str, Any]:
    meta = await oidc_discovery()
    async with httpx.AsyncClient(timeout=15) as client:
        token_resp = await client.post(meta["token_endpoint"], data={
            "grant_type": "authorization_code",
            "code": code,
            "redirect_uri": redirect_uri,
            "client_id": settings.oidc_client_id,
            "client_secret": settings.oidc_client_secret,
        }, headers={"Accept": "application/json"})
        if token_resp.status_code != 200:
            raise OIDCError(f"Token exchange failed ({token_resp.status_code}): {token_resp.text[:200]}")
        access = token_resp.json().get("access_token")
        if not access:
            raise OIDCError("Token response contained no access_token.")

        info_resp = await client.get(meta["userinfo_endpoint"],
                                     headers={"Authorization": f"Bearer {access}"})
        if info_resp.status_code != 200:
            raise OIDCError(f"userinfo failed ({info_resp.status_code}): {info_resp.text[:200]}")
        return info_resp.json()


def upsert_oidc_user(db: Session, claims: dict[str, Any]) -> User:
    """Find or create the local record for an SSO identity.

    Matching is on `sub` first so a changed email still lands on the same user;
    an existing local account with the same address is adopted rather than
    duplicated.
    """
    sub = str(claims.get("sub") or "").strip()
    email = str(claims.get("email") or "").strip().lower()
    if not sub and not email:
        raise OIDCError("SSO returned neither sub nor email.")

    groups = claims.get(settings.oidc_groups_claim) or []
    if isinstance(groups, str):
        groups = [groups]
    is_admin = settings.oidc_admin_group in {str(g) for g in groups}

    user = None
    if sub:
        user = db.query(User).filter(User.oidc_sub == sub).one_or_none()
    if user is None and email:
        user = db.query(User).filter(User.email == email).one_or_none()

    if user is None:
        user = User(email=email or f"{sub}@sso.local", provider=AuthProvider.OIDC)
        db.add(user)

    user.oidc_sub = sub or user.oidc_sub
    user.provider = AuthProvider.OIDC
    user.name = str(claims.get("name") or claims.get("preferred_username") or user.name or email)
    if email:
        user.email = email
    user.is_admin = is_admin
    user.is_active = True
    user.last_login = dt.datetime.now(dt.timezone.utc)
    if not user.signoff_name:
        user.signoff_name = user.name
    db.commit()
    return user