115 lines
3.3 KiB
Python
115 lines
3.3 KiB
Python
"""Keycloak JWT authentication middleware."""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
from dotenv import load_dotenv
|
|
from fastapi import Depends, HTTPException, Request, status
|
|
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
|
from jose import JWTError, jwt
|
|
from jose.constants import Algorithms
|
|
|
|
load_dotenv(Path(__file__).parent.parent / ".env", override=True)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
security = HTTPBearer(auto_error=False)
|
|
|
|
KEYCLOAK_SERVER = os.getenv("KEYCLOAK_SERVER_URL", "https://keycloak.plfai.cn")
|
|
KEYCLOAK_REALM = os.getenv("KEYCLOAK_REALM", "plfai")
|
|
KEYCLOAK_AUDIENCE = os.getenv("KEYCLOAK_AUDIENCE", "account")
|
|
|
|
OIDC_CONFIG_URL = f"{KEYCLOAK_SERVER}/realms/{KEYCLOAK_REALM}/.well-known/openid-configuration"
|
|
|
|
# Cached OIDC config + JWKS
|
|
_oidc_config: Optional[dict] = None
|
|
_jwks: Optional[dict] = None
|
|
|
|
|
|
async def _fetch_oidc_config() -> tuple[dict, dict]:
|
|
"""Fetch OIDC config and JWKS, caching them."""
|
|
global _oidc_config, _jwks
|
|
|
|
if _oidc_config is None:
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.get(OIDC_CONFIG_URL)
|
|
resp.raise_for_status()
|
|
_oidc_config = resp.json()
|
|
|
|
jwks_resp = await client.get(_oidc_config["jwks_uri"])
|
|
jwks_resp.raise_for_status()
|
|
_jwks = jwks_resp.json()
|
|
|
|
return _oidc_config, _jwks # type: ignore[return-value]
|
|
|
|
|
|
async def get_current_user(
|
|
request: Request,
|
|
credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
|
|
) -> dict:
|
|
"""Validate Bearer token and return user claims. Raises 401 if invalid."""
|
|
if credentials is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Missing Authorization header",
|
|
)
|
|
|
|
token = credentials.credentials
|
|
|
|
try:
|
|
_, jwks = await _fetch_oidc_config()
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch OIDC config: {e}")
|
|
raise HTTPException(
|
|
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
detail="Auth service unavailable",
|
|
)
|
|
|
|
# Get kid from token header
|
|
try:
|
|
unverified = jwt.get_unverified_header(token)
|
|
except JWTError:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Invalid token header",
|
|
)
|
|
|
|
kid = unverified.get("kid")
|
|
if not kid:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Missing kid in token",
|
|
)
|
|
|
|
# Find matching public key
|
|
key = None
|
|
for k in jwks.get("keys", []):
|
|
if k.get("kid") == kid:
|
|
key = k
|
|
break
|
|
|
|
if key is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="Unknown signing key",
|
|
)
|
|
|
|
# Verify signature + expiry
|
|
try:
|
|
payload = jwt.decode(
|
|
token,
|
|
key,
|
|
algorithms=[Algorithms.RS256],
|
|
audience=KEYCLOAK_AUDIENCE,
|
|
options={"verify_exp": True},
|
|
)
|
|
return payload
|
|
except JWTError as e:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail=f"Invalid token: {e}",
|
|
)
|