File size: 3,824 Bytes
db4ba8d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
621eb30
 
 
 
 
 
db4ba8d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
"""
TradeFlow AI — Keycloak 26 JWT Authentication (T-008)

PRD Invariant #4: Keycloak 26 is the SOLE auth provider. No Supabase Auth.
Validates JWTs via JWKS endpoint with a 5-minute cache.
"""

from __future__ import annotations

import time
from typing import Any

import httpx
from fastapi import HTTPException, status
from jose import JWTError, jwk, jwt

from ..config import settings

# JWKS TTL — 5 minutes (SDD §4.1)
_JWKS_TTL_SECONDS = 300


class JWKSCache:
    """Thread-safe JWKS cache with 5-minute TTL."""

    def __init__(self) -> None:
        self._keys: dict[str, Any] = {}
        self._fetched_at: float = 0.0

    def _is_stale(self) -> bool:
        return time.monotonic() - self._fetched_at > _JWKS_TTL_SECONDS

    async def get_keys(self) -> dict[str, Any]:
        if not self._keys or self._is_stale():
            await self._refresh()
        return self._keys

    async def _refresh(self) -> None:
        async with httpx.AsyncClient(timeout=10.0) as client:
            resp = await client.get(settings.KEYCLOAK_JWKS_URL)
            resp.raise_for_status()
            jwks_data = resp.json()

        # Build kid → key mapping
        self._keys = {}
        for key_data in jwks_data.get("keys", []):
            kid = key_data.get("kid")
            if kid:
                self._keys[kid] = jwk.construct(key_data)

        self._fetched_at = time.monotonic()


_jwks_cache = JWKSCache()


async def verify_keycloak_token(token: str) -> dict[str, Any]:
    """
    Verify a Keycloak JWT token.

    Returns the decoded payload (claims) on success.
    Raises HTTP 401 on any failure.
    """
    if settings.DISABLE_AUTH:
        return {
            "sub": "demo-bypass-user",
            "realm_access": {"roles": ["operator", "admin", "supervisor", "sme"]},
        }

    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )

    try:
        # Step 1: Decode header to get kid without signature verification
        unverified_header = jwt.get_unverified_header(token)
    except JWTError:
        raise credentials_exception

    kid = unverified_header.get("kid")
    if not kid:
        raise credentials_exception

    # Step 2: Get the signing key from JWKS cache
    try:
        keys = await _jwks_cache.get_keys()
    except Exception:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Auth service temporarily unavailable",
        )

    signing_key = keys.get(kid)
    if not signing_key:
        # Key not found — JWKS may have rotated, force refresh
        await _jwks_cache._refresh()
        keys = await _jwks_cache.get_keys()
        signing_key = keys.get(kid)
        if not signing_key:
            raise credentials_exception

    # Step 3: Verify signature + claims
    try:
        payload = jwt.decode(
            token,
            signing_key,
            algorithms=["RS256"],
            audience=settings.KEYCLOAK_CLIENT_ID,
            issuer=settings.KEYCLOAK_ISSUER,
            options={"verify_exp": True},
        )
    except JWTError as e:
        raise credentials_exception from e

    return payload


def extract_roles(payload: dict[str, Any]) -> list[str]:
    """Extract realm-level roles from a decoded Keycloak token."""
    realm_access = payload.get("realm_access", {})
    return realm_access.get("roles", [])


def extract_user_id(payload: dict[str, Any]) -> str:
    """Extract the user UUID (Keycloak sub claim)."""
    sub = payload.get("sub")
    if not sub:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Token missing subject claim",
        )
    return sub