Coverage for src/lexigram/auth/authn/_jwt_lifecycle.py: 84%
153 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:58 +0800
1"""Token verification and lifecycle mixin for :class:`~lexigram.auth.authn.jwt.JWTTokenManager`.
3This module is an internal implementation detail; import
4:class:`~lexigram.auth.authn.jwt.JWTTokenManager` directly.
5"""
7from __future__ import annotations
9from collections import OrderedDict
10import hashlib
11from typing import TYPE_CHECKING, Any
13import jwt
15from lexigram.auth.authn._binding import TokenBindingConfig, verify_binding
16from lexigram.auth.authn.blacklist import JWTBlacklist
17from lexigram.auth.models import AuthToken
18from lexigram.auth.models.user import User
19from lexigram.contracts.auth.exceptions import TokenError as ContractsTokenError
21if TYPE_CHECKING:
22 from lexigram.auth.types import TokenPair
23 from lexigram.contracts.audit import AuditLoggerProtocol
24 from lexigram.contracts.auth.token import VerifiedToken
25 from lexigram.contracts.core import HookRegistryProtocol
26 from lexigram.logging import LoggerProtocol as Logger
27 from lexigram.result import Result
30class _JWTLifecycleMixin:
31 """Mixin providing JWT token verification and lifecycle methods for :class:`JWTTokenManager`.
33 All public attributes referenced here are initialised by
34 ``JWTTokenManager.__init__``; they are declared below as class-level
35 annotations solely to satisfy static type checkers.
36 """
38 # ── Attributes set by JWTTokenManager.__init__ ───────────────────────────
39 algorithm: str
40 access_expiration_hours: int
41 refresh_expiration_days: int
42 _required_audience: str | None
43 _binding_config: TokenBindingConfig | None
44 _blacklist_mgr: JWTBlacklist
45 _verification_cache: OrderedDict[str, str]
46 _verified_by_key: dict[str, str]
47 logger: Logger
48 _audit_logger: AuditLoggerProtocol | None
49 _hooks: HookRegistryProtocol | None
51 @property
52 def keys(self) -> dict[str, Any]: # pragma: no cover
53 """Live key material — provided by JWTTokenManager."""
54 raise NotImplementedError
56 @property
57 def current_key_id(self) -> str: # pragma: no cover
58 """Active signing key ID — provided by JWTTokenManager."""
59 raise NotImplementedError
61 def _get_verification_key(self, kid: str) -> str: # pragma: no cover
62 """Return raw verification key — provided by JWTTokenManager."""
63 raise NotImplementedError
65 def create_token_pair( # pragma: no cover
66 self,
67 user: User,
68 additional_claims: dict[str, Any] | None = None,
69 binding_context: dict[str, str] | None = None,
70 ) -> AuthToken:
71 """Create token pair — provided by _JWTCreationMixin."""
72 raise NotImplementedError
74 # ─────────────────────────────────────────────────────────────────────────
76 async def _emit_action(self, hook_name: str, payload: object) -> None:
77 """Emit a token lifecycle hook when a registry is available."""
78 if self._hooks is None:
79 return
81 await self._hooks.call_action(hook_name, payload=payload)
83 async def refresh_token(
84 self, refresh_token: str
85 ) -> Result[AuthToken, ContractsTokenError]:
86 """Refresh an access token using a refresh token.
88 Delegates to :meth:`refresh_access_token`. Satisfies the updated
89 ``TokenManagerProtocol`` protocol from ``lexigram.contracts.auth.token``.
91 Args:
92 refresh_token: The refresh token string.
94 Returns:
95 ``Ok(AuthToken)`` if the refresh token is valid, or
96 ``Err(TokenError)`` for expected domain failures
97 (expired, blacklisted, invalid).
98 """
99 from lexigram.auth.exceptions import (
100 BlacklistedTokenError,
101 TokenBlacklistedError,
102 )
103 from lexigram.auth.exceptions import TokenError as AuthTokenError
104 from lexigram.result import Err, Ok
106 try:
107 return Ok(await self.refresh_access_token(refresh_token))
108 except BlacklistedTokenError as e:
109 return Err(TokenBlacklistedError(str(e))) # type: ignore[arg-type]
110 except AuthTokenError as e:
111 return Err(ContractsTokenError(str(e)))
113 async def refresh_with_rotation(
114 self, refresh_token: str
115 ) -> Result[TokenPair, ContractsTokenError]:
116 """Rotate a refresh token and return a new access + refresh token pair.
118 Implements one-time-use refresh token rotation (RTR). The incoming
119 refresh token is verified, immediately blacklisted, and a fresh token
120 pair is issued. If the token is already blacklisted (reuse-attack
121 vector), all sessions for the affected user are revoked and the call
122 returns ``Err``.
124 Args:
125 refresh_token: The current refresh token JWT string to exchange.
127 Returns:
128 ``Ok(TokenPair)`` with fresh ``access`` and ``refresh`` tokens on
129 success, or ``Err(ContractsTokenError)`` for any expected domain
130 failure (expired, invalid signature, already revoked).
132 Raises:
133 RuntimeError: On cache / infrastructure failures.
134 OSError: On network-level failures.
135 ConnectionError: On connection failures.
136 """
137 from lexigram.auth.exceptions import (
138 BlacklistedTokenError,
139 TokenBlacklistedError,
140 )
141 from lexigram.auth.exceptions import TokenError as AuthTokenError
142 from lexigram.auth.types import TokenPair
143 from lexigram.result import Err, Ok
145 try:
146 auth_token = await self.refresh_access_token(refresh_token)
147 return Ok(
148 TokenPair(
149 access=auth_token.token,
150 refresh=auth_token.refresh_token or "",
151 )
152 )
153 except BlacklistedTokenError as e:
154 return Err(TokenBlacklistedError(str(e))) # type: ignore[arg-type]
155 except AuthTokenError as e:
156 return Err(ContractsTokenError(str(e)))
158 async def verify_token(
159 self,
160 token: str,
161 token_type: str = "access", # noqa: S107 # not a password; identifies token kind
162 expected_audience: str | None = None,
163 required_scope: str | None = None,
164 binding_context: dict[str, str] | None = None,
165 *,
166 allow_missing_audience: bool = False,
167 ) -> Result[VerifiedToken, ContractsTokenError]:
168 """Verify and decode a JWT token with support for multiple keys.
170 Returns a ``Result`` rather than raising domain exceptions. Only
171 infrastructure failures (cache down, network errors) are still raised
172 as exceptions so the event loop can propagate them properly.
174 Args:
175 token: The JWT token to verify.
176 token_type: Expected token type ("access" or "refresh").
177 expected_audience: Expected audience claim ("aud") for validation.
178 Takes precedence over ``required_audience`` set on the manager.
179 required_scope: Required scope claim for validation.
180 binding_context: Optional request context dict for token binding.
181 allow_missing_audience: When ``True``, audience validation is skipped
182 entirely regardless of ``required_audience`` on the manager.
183 Use only for trusted internal service paths that explicitly opt
184 out of audience enforcement.
186 Returns:
187 ``Ok(VerifiedToken)`` if the token is valid and not revoked.
188 ``Err(TokenError)`` for expected domain failures (expired,
189 blacklisted, invalid signature, wrong type, audience mismatch, etc.).
191 Raises:
192 RuntimeError: If the cache backend is unavailable (infrastructure).
193 OSError: On network-level failures (infrastructure).
194 ConnectionError: On connection failures (infrastructure).
195 """
196 from lexigram.auth.exceptions import TokenAudienceError, TokenInvalidError
197 from lexigram.auth.exceptions import (
198 TokenBlacklistedError as ContractsBlacklistedError,
199 )
200 from lexigram.auth.exceptions import (
201 TokenExpiredError as ContractsExpiredError,
202 )
203 from lexigram.contracts.auth.token import VerifiedToken
204 from lexigram.result import Err, Ok
206 try:
207 # Compute a short token hash for the verification cache.
208 # 16 hex chars (64 bits) is sufficient for an in-process lookup;
209 # the full token is never stored.
210 token_hash = hashlib.sha256(token.encode()).hexdigest()
212 # Get key ID from header
213 header = jwt.get_unverified_header(token)
214 kid = header.get("kid", self.current_key_id)
216 # Check if key exists
217 if kid not in self.keys:
218 return Err(TokenInvalidError("Unknown signing key ID")) # type: ignore[arg-type]
220 # Fast path: if we have already verified this exact token recently,
221 # skip the kid-based lookup and go straight to the cached key_id.
222 # Fall back to the per-kid cache, then the kid header itself.
223 cached_key_id = self._verification_cache.get(token_hash)
224 if cached_key_id and cached_key_id in self.keys:
225 effective_key_id = cached_key_id
226 self._verification_cache.move_to_end(
227 token_hash
228 ) # mark as recently used
229 else:
230 # Try the key that last successfully verified this kid (fast path),
231 # fall back to kid itself if the cached entry is absent or stale.
232 effective_key_id = self._verified_by_key.get(kid, kid)
233 if effective_key_id not in self.keys:
234 effective_key_id = kid
236 # Verify with correct key
237 verification_key = self._get_verification_key(effective_key_id)
239 # Resolve the effective audience for this call:
240 # 1. An explicit call-site audience always wins.
241 # 2. Fall back to the manager-level required_audience.
242 # 3. When allow_missing_audience=True, skip the check entirely.
243 if allow_missing_audience:
244 effective_audience: str | None = None
245 decode_options: dict = {"verify_aud": False}
246 else:
247 effective_audience = expected_audience or self._required_audience
248 decode_options = {"verify_aud": effective_audience is not None}
249 if effective_audience is None:
250 import os
252 if os.getenv("LEX_ENV", "development") == "production":
253 self.logger.warning(
254 "jwt_verification_without_audience",
255 token_type=token_type,
256 )
258 payload = jwt.decode(
259 token,
260 verification_key,
261 algorithms=[self.algorithm],
262 audience=effective_audience,
263 options=decode_options, # type: ignore[arg-type]
264 )
266 # Record which key_id successfully verified this kid and this token.
267 self._verified_by_key[kid] = effective_key_id
268 if len(self._verification_cache) >= 1000:
269 self._verification_cache.popitem(last=False) # evict LRU (oldest)
270 self._verification_cache[token_hash] = effective_key_id
272 # Check token type — only enforce if the payload explicitly declares one.
273 # External JWTs (e.g. NextAuth HS256 tokens) omit the "type" claim; rejecting
274 # them here would silently break all external-JWT auth flows.
275 declared_type = payload.get("type")
276 if declared_type is not None and declared_type != token_type:
277 return Err(
278 TokenInvalidError(f"Invalid token type: expected {token_type}") # type: ignore[arg-type]
279 )
281 # Validate audience explicitly if jwt.decode didn't handle it fully
282 if effective_audience:
283 token_aud = payload.get("aud")
284 if isinstance(token_aud, list):
285 if effective_audience not in token_aud:
286 return Err(
287 TokenAudienceError( # type: ignore[arg-type]
288 f"Invalid audience: {effective_audience}"
289 )
290 )
291 elif token_aud != effective_audience:
292 return Err(
293 TokenAudienceError(f"Invalid audience: {effective_audience}") # type: ignore[arg-type]
294 )
296 # Validate scope if required
297 if required_scope:
298 token_scope = payload.get("scope", "")
299 if required_scope not in token_scope.split():
300 return Err(
301 TokenInvalidError(f"Required scope missing: {required_scope}") # type: ignore[arg-type]
302 )
304 # Verify client binding if configured
305 if self._binding_config and binding_context is not None:
306 if not verify_binding(self._binding_config, payload, binding_context):
307 return Err(TokenInvalidError("Token binding mismatch")) # type: ignore[arg-type]
309 # Check if token is blacklisted — may raise RuntimeError (infra)
310 if await self._is_token_blacklisted(token):
311 return Err(ContractsBlacklistedError("Token has been revoked")) # type: ignore[arg-type]
313 from datetime import UTC
315 # Surface application-defined claims; registered + known claims
316 # are already mapped to typed fields above.
317 _known = {
318 "sub",
319 "email",
320 "name",
321 "roles",
322 "permissions",
323 "exp",
324 "iat",
325 "nbf",
326 "iss",
327 "jti",
328 "aud",
329 "type",
330 "scope",
331 }
332 extra_claims = {
333 key: value
334 for key, value in payload.items()
335 if key not in _known and not key.startswith("_")
336 }
338 return Ok(
339 VerifiedToken(
340 user_id=payload.get("sub", ""),
341 email=payload.get("email", ""),
342 name=payload.get("name", ""),
343 roles=payload.get("roles", []),
344 permissions=payload.get("permissions", []),
345 expires_at=__import__("datetime").datetime.fromtimestamp(
346 payload.get("exp", 0), tz=UTC
347 ),
348 key_id=kid,
349 token_type=token_type,
350 audience=effective_audience,
351 extra_claims=extra_claims,
352 )
353 )
355 except jwt.ExpiredSignatureError:
356 return Err(ContractsExpiredError("Token has expired")) # type: ignore[arg-type]
357 except jwt.InvalidAudienceError:
358 return Err(TokenAudienceError("Invalid audience")) # type: ignore[arg-type]
359 except (jwt.InvalidTokenError, ValueError) as e:
360 return Err(TokenInvalidError(f"Invalid token: {e}")) # type: ignore[arg-type]
362 async def logout(self, token: str) -> Result[None, ContractsTokenError]:
363 """Invalidate a token by adding it to the blacklist.
365 When a ``cache_service`` is configured the token hash is stored in
366 the cache with key ``jwt:blacklist:{sha256_hex}`` and a TTL equal to
367 the token's remaining lifetime. When no cache is configured the hash
368 is stored in an in-process set (cleared on process restart).
370 Args:
371 token: The JWT token to invalidate.
373 Returns:
374 ``Ok(None)`` on success or if the token was already expired.
375 ``Err(TokenError)`` if the cache write failed.
377 Raises:
378 RuntimeError: On cache-level infrastructure failures.
379 OSError: On network-level failures.
380 ConnectionError: On connection failures.
381 """
382 from lexigram.auth.hooks import AuthTokenRevokedHook
384 result = await self._blacklist_mgr.revoke(token)
385 if result.is_ok():
386 try:
387 payload = jwt.decode(
388 token,
389 options={
390 "verify_signature": False,
391 "verify_exp": False,
392 "verify_aud": False,
393 },
394 )
395 except (jwt.DecodeError, jwt.InvalidTokenError, ValueError, TypeError):
396 self.logger.warning("jwt_logout_hook_payload_decode_failed")
397 else:
398 await self._emit_action(
399 "auth.logout",
400 AuthTokenRevokedHook(
401 user_id=str(payload.get("sub", "")),
402 token_type=str(payload.get("type", "access")),
403 ),
404 )
405 return result
407 async def logout_all_user_tokens(
408 self, user_id: str
409 ) -> Result[None, ContractsTokenError]:
410 """Invalidate all tokens for a user by writing a user-level blacklist entry.
412 Args:
413 user_id: The user ID to invalidate all tokens for.
415 Returns:
416 ``Ok(None)`` on success. ``Err(TokenError)`` if the cache write failed.
418 Raises:
419 RuntimeError: If no cache backend is configured.
420 """
421 return await self._blacklist_mgr.revoke_all_for_user(user_id)
423 async def _is_token_blacklisted(self, token: str) -> bool:
424 """Check if a token is blacklisted.
426 Args:
427 token: The JWT token to check.
429 Returns:
430 ``True`` if the token should be rejected, ``False`` otherwise.
431 """
432 return await self._blacklist_mgr.is_blacklisted(token)
434 async def refresh_access_token(self, refresh_token: str) -> AuthToken:
435 """Create new access token from refresh token with rotation.
437 Implements one-time-use refresh tokens (RTR). The used refresh token is
438 blacklisted and a new refresh token is issued.
440 If a blacklisted refresh token is presented, it indicates a potential
441 reuse attack. In this case, all tokens for the user are invalidated
442 immediately to prevent further compromise.
444 Raises:
445 TokenBlacklistedError: If token reuse is detected (domain failure).
446 TokenError: For other domain failures (expired, invalid).
447 RuntimeError: On infrastructure failures (cache, network).
448 """
449 from lexigram.auth.exceptions import (
450 BlacklistedTokenError,
451 TokenBlacklistedError,
452 )
453 from lexigram.auth.exceptions import TokenError as AuthTokenError
455 result = await self.verify_token(refresh_token, "refresh")
456 if result.is_err():
457 error = result.unwrap_err()
458 if isinstance(error, TokenBlacklistedError):
459 # TOKEN REUSE DETECTED!
460 # If we try to use a blacklisted refresh token, someone might have stolen it.
461 # Invalidate EVERYTHING for this user.
462 try:
463 unverified_payload = jwt.decode(
464 refresh_token, options={"verify_signature": False}
465 )
466 user_id = unverified_payload.get("sub")
467 if user_id:
468 self.logger.warning(
469 "Refresh token reuse detected for user %s. Revoking all sessions.",
470 user_id,
471 )
472 await self.logout_all_user_tokens(user_id)
473 except (RuntimeError, OSError, ConnectionError):
474 self.logger.exception(
475 "Failed to revoke user tokens after reuse detection",
476 )
477 except (jwt.DecodeError, jwt.InvalidTokenError, ValueError, KeyError):
478 self.logger.exception(
479 "Unexpected error during session revocation",
480 )
481 raise BlacklistedTokenError(
482 "Refresh token reuse detected. All sessions revoked.",
483 ) from None
484 # Other domain errors (expired, invalid)
485 raise AuthTokenError(str(error)) from error
487 verified = result.unwrap()
489 # Blacklist the used refresh token (rotation)
490 await self.logout(refresh_token)
492 # Create new token pair from token claims
493 user = User(
494 user_id=verified.user_id,
495 name=verified.name,
496 email=verified.email,
497 roles=verified.roles,
498 permissions=verified.permissions,
499 )
501 from lexigram.auth.hooks import AuthTokenRefreshedHook
503 token_pair = self.create_token_pair(user)
504 await self._emit_action(
505 "token.refreshed",
506 AuthTokenRefreshedHook(user_id=verified.user_id, token_type="access"), # noqa: S106 # token KIND string, not a credential
507 )
508 return token_pair
510 async def get_user_from_token(
511 self, token: str
512 ) -> Result[VerifiedToken, ContractsTokenError]:
513 """Extract user information from access token.
515 Returns:
516 ``Ok(VerifiedToken)`` if the token is a valid access token, or
517 ``Err(TokenError)`` for expected domain failures.
518 """
519 return await self.verify_token(token, "access")