ai-chat/backend/app/auth.py

import logging
from functools import lru_cache
from typing import Any, Dict, Optional
from urllib.parse import urljoin

import httpx
from fastapi import HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwk, jwt
from pydantic import BaseModel, Field, HttpUrl
from pydantic_settings import BaseSettings

logger = logging.getLogger(__name__)

security = HTTPBearer(auto_error=False)

class Settings(BaseSettings):
    # OIDC issuer (discovery URL root)
    oidc_issuer: HttpUrl = Field(..., env="OIDC_ISSUER")
    # Optional audience: if set the backend will validate the token 'aud' claim.
    # If omitted, audience validation is skipped (use with caution).
    oidc_audience: Optional[str] = Field(None, env="OIDC_AUDIENCE")
    oidc_client_id: str = Field(..., env="OIDC_CLIENT_ID")
    oidc_client_secret: Optional[str] = Field(None, env="OIDC_CLIENT_SECRET")
    oidc_verify_iss: bool = Field(True, env="OIDC_VERIFY_ISS")

    class Config:
        env_file = ".env"

settings = Settings()

class TokenClaims(BaseModel):
    iss: str
    sub: str
    aud: Any
    exp: int
    iat: Optional[int]
    azp: Optional[str]
    email: Optional[str]
    preferred_username: Optional[str]

@lru_cache(maxsize=1)
def get_discovery() -> Dict[str, Any]:
    issuer = str(settings.oidc_issuer).rstrip("/")
    discovery_url = urljoin(issuer + "/", ".well-known/openid-configuration")
    with httpx.Client(timeout=10.0) as client:
        response = client.get(discovery_url)
        if response.status_code != 200:
            raise RuntimeError(f"Unable to load OIDC discovery document from {discovery_url}")
        return response.json()

@lru_cache(maxsize=1)
def get_jwks() -> Dict[str, Any]:
    discovery = get_discovery()
    jwks_uri = discovery.get("jwks_uri")
    if not jwks_uri:
        raise RuntimeError("OIDC discovery document does not contain jwks_uri")
    with httpx.Client(timeout=10.0) as client:
        response = client.get(jwks_uri)
        if response.status_code != 200:
            raise RuntimeError(f"Unable to load JWKS from {jwks_uri}")
        return response.json()

def get_signing_key(token: str) -> Dict[str, Any]:
    try:
        headers = jwt.get_unverified_header(token)
    except JWTError as exc:
        raise HTTPException(status_code=401, detail="Invalid JWT header") from exc

    kid = headers.get("kid")
    if not kid:
        raise HTTPException(status_code=401, detail="JWT missing kid header")

    jwks = get_jwks().get("keys", [])
    for key in jwks:
        if key.get("kid") == kid:
            return key

    raise HTTPException(status_code=401, detail="Unable to find matching JWKS key")

@lru_cache(maxsize=1)
def get_introspection_endpoint() -> Optional[str]:
    discovery = get_discovery()
    return discovery.get("introspection_endpoint")

def introspect_token(token: str) -> Dict[str, Any]:
    introspection_endpoint = get_introspection_endpoint()
    if not introspection_endpoint:
        raise HTTPException(status_code=401, detail="Unable to introspect token")

    data = {"token": token, "token_type_hint": "access_token"}
    auth = None
    if settings.oidc_client_secret:
        auth = (settings.oidc_client_id, settings.oidc_client_secret)
    else:
        data["client_id"] = settings.oidc_client_id

    with httpx.Client(timeout=10.0) as client:
        response = client.post(introspection_endpoint, data=data, headers={"Accept": "application/json"}, auth=auth)

    if response.status_code != 200:
        raise HTTPException(status_code=401, detail="Unable to introspect token")

    introspection = response.json()
    if not introspection.get("active"):
        raise HTTPException(status_code=401, detail="Invalid or expired token")

    discovery = get_discovery()
    if "iss" not in introspection:
        introspection["iss"] = discovery.get("issuer")

    return introspection

async def validate_token(credentials: HTTPAuthorizationCredentials = Security(security)) -> TokenClaims:
    if not credentials or credentials.scheme.lower() != "bearer":
        raise HTTPException(status_code=401, detail="Missing Bearer authorization header")

    token = credentials.credentials
    try:
        key = get_signing_key(token)
        public_key = jwk.construct(key)
        discovery = get_discovery()
        # Only pass `audience` to the decoder if configured. Some providers
        # (or local development setups) may not include the aud claim in a
        # way that matches your API identifier; in that case leave
        # `OIDC_AUDIENCE` unset and the audience check will be skipped.
        jwt_kwargs = {
            "algorithms": [key.get("alg", "RS256")],
            "issuer": discovery.get("issuer") if settings.oidc_verify_iss else None,
        }
        if settings.oidc_audience:
            jwt_kwargs["audience"] = settings.oidc_audience

        verified = jwt.decode(token, public_key, **jwt_kwargs)
    except JWTError:
        logger.exception("JWT validation failed, attempting introspection")
        verified = introspect_token(token)

    try:
        return TokenClaims(**verified)
    except Exception as exc:
        raise HTTPException(status_code=401, detail="Unable to parse token claims") from exc

    try:
        return TokenClaims(**verified)
    except Exception as exc:
        raise HTTPException(status_code=401, detail="Unable to parse token claims") from exc