Coverage for src/lexigram/auth/authn/jwt.py: 88%

92 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-26 00:58 +0800

1"""JWT token management for authentication. 

2 

3Provides :class:`JWTTokenManager` for JWT creation, validation, and refresh. 

4Key lifecycle is delegated to :mod:`key_rotation` (:class:`JWTKeyStore`) and 

5persistent revocation to :mod:`revocation` (:class:`PersistentTokenRevocationStore`). 

6 

7Example:: 

8 

9 manager = JWTTokenManager( 

10 current_key_id="key-1", 

11 keys={"key-1": SecretStr("secret-key-at-least-32-chars")}, 

12 ) 

13 token = await manager.create_access_token(user_id="123") 

14 result = await manager.verify_token(token) 

15""" 

16 

17from __future__ import annotations 

18 

19from collections import OrderedDict 

20import os 

21from typing import TYPE_CHECKING, Any 

22 

23from lexigram.auth import constants as const 

24from lexigram.auth.authn._binding import TokenBindingConfig 

25from lexigram.auth.authn._jwt_creation import _JWTCreationMixin 

26from lexigram.auth.authn._jwt_lifecycle import _JWTLifecycleMixin 

27from lexigram.auth.authn._key_utils import normalize_jwt_keys 

28from lexigram.auth.authn.blacklist import JWTBlacklist 

29from lexigram.auth.authn.key_rotation import JWTKeyStore 

30from lexigram.contracts.exceptions import ConfigurationError 

31from lexigram.logging import LoggerProtocol as Logger 

32from lexigram.logging import get_logger 

33from lexigram.validation import SecretStr 

34 

35_logger = get_logger(__name__) 

36 

37if TYPE_CHECKING: 

38 from lexigram.contracts.audit import AuditLoggerProtocol 

39 from lexigram.contracts.auth.exceptions import TokenError as ContractsTokenError 

40 from lexigram.contracts.auth.token import VerifiedToken 

41 from lexigram.contracts.core import HookRegistryProtocol 

42 from lexigram.contracts.core.identity import IdGeneratorProtocol 

43 from lexigram.contracts.infra.cache import CacheBackendProtocol 

44 from lexigram.result import Result 

45 

46 

47class JWTTokenManager(_JWTCreationMixin, _JWTLifecycleMixin): 

48 """JWT token management with key rotation support. 

49 

50 The JWTTokenManager handles creation, validation, and renewal of JWT 

51 tokens for user authentication. It supports multiple signing keys 

52 for seamless key rotation without invalidating existing tokens. 

53 

54 Attributes: 

55 current_key_id: ID of the currently active signing key. 

56 keys: Dictionary of key_id -> key material. 

57 algorithm: JWT signing algorithm (HS256, RS256, etc.). 

58 access_expiration_hours: Expiration time for access tokens. 

59 refresh_expiration_days: Expiration time for refresh tokens. 

60 cache_service: Optional cache backend for token validation caching. 

61 rotation_interval: Interval in seconds between key rotations. 

62 

63 Example: 

64 Basic token operations:: 

65 

66 manager = JWTTokenManager( 

67 current_key_id="v1", 

68 keys={"v1": SecretStr("secure-secret-key")}, 

69 ) 

70 

71 # Create access token 

72 token = await manager.create_access_token(user_id="user123") 

73 

74 # Verify and decode 

75 payload = await manager.verify_token(token) 

76 

77 Note: 

78 In production, secrets should not be hardcoded. Use environment 

79 variables or a secrets management service. 

80 """ 

81 

82 def __init__( 

83 self, 

84 current_key_id: str, 

85 keys: dict[str, str | SecretStr | dict[str, str | SecretStr]] | None = None, 

86 algorithm: str = const.DEFAULT_TOKEN_ALGORITHM, 

87 access_expiration_hours: int = 24, 

88 refresh_expiration_days: int = 30, 

89 cache_service: CacheBackendProtocol | None = None, 

90 rotation_interval_days: int = 90, 

91 grace_period_seconds: int = const.DEFAULT_JWT_KEY_ROTATION_GRACE_PERIOD_SECONDS, 

92 logger: Logger | None = None, 

93 *, 

94 audit_logger: AuditLoggerProtocol | None = None, 

95 binding_config: TokenBindingConfig | None = None, 

96 required_audience: str | None = None, 

97 ids: IdGeneratorProtocol | None = None, 

98 ) -> None: 

99 """Initialize the JWT token manager. 

100 

101 Args: 

102 current_key_id: The ID of the currently active signing key. 

103 keys: Dictionary mapping key IDs to key material. If None, 

104 current_key_id is treated as a single secret. 

105 algorithm: JWT signing algorithm (e.g., "HS256", "RS256"). 

106 access_expiration_hours: Hours until access tokens expire. 

107 refresh_expiration_days: Days until refresh tokens expire. 

108 cache_service: Optional cache backend for token validation. 

109 rotation_interval_days: Days between key rotations. 

110 grace_period_seconds: Seconds after key rotation during which tokens 

111 signed by the outgoing key remain valid. Defaults to 3600 (1 hour). 

112 Set to 0 to invalidate all old-key tokens immediately on rotation. 

113 audit_logger: Optional :class:`~lexigram.contracts.audit.AuditLoggerProtocol` 

114 used to record token revocation events. When not provided, 

115 no audit entries are written. 

116 binding_config: Optional :class:`TokenBindingConfig` for opt-in client 

117 binding. When set, tokens embed a ``bind`` claim containing a 

118 SHA-256 hash of active binding factors (IP, fingerprint). Tokens 

119 issued without a ``bind`` claim continue to verify successfully so 

120 that binding can be enabled incrementally. 

121 required_audience: When set, every call to ``verify_token`` will 

122 enforce that the token's ``aud`` claim matches this value. 

123 Pass ``allow_missing_audience=True`` to ``verify_token`` to 

124 bypass the check on a per-call basis for trusted internal paths. 

125 """ 

126 if keys is None: 

127 # Single secret mode for backward compatibility and simple usage 

128 raw = current_key_id 

129 _resolved_key_id = "default" 

130 _resolved_keys = normalize_jwt_keys({"default": raw}) 

131 else: 

132 if not keys: 

133 raise ValueError("keys cannot be empty") 

134 _resolved_key_id = current_key_id 

135 _resolved_keys = normalize_jwt_keys(keys) 

136 

137 # Ensure the current key exists in the keys dictionary 

138 if _resolved_key_id not in _resolved_keys: 

139 raise ValueError( 

140 f"current_key_id '{_resolved_key_id}' not found in keys", 

141 ) 

142 

143 # Defensive GuardProtocol: Ensure secret is not a known weak default in production 

144 env = os.getenv("LEX_ENV", "development").lower() 

145 if env in ("production", "staging"): 

146 weak_secrets = ["change-me-in-production", "secret", "password", "123456"] 

147 for key_val in _resolved_keys.values(): 

148 if isinstance(key_val, dict): 

149 # Asymmetric keys — skip weak-key check (PEM keys are always long) 

150 continue 

151 secret_str = key_val.get_secret_value() 

152 if any(w in secret_str.lower() for w in weak_secrets) or ( 

153 len(secret_str) < 32 

154 ): 

155 raise ConfigurationError( 

156 "Insecure JWT secret key detected in production environment", 

157 ) 

158 

159 # Delegate all key lifecycle management to JWTKeyStore 

160 self._key_store = JWTKeyStore( 

161 current_key_id=_resolved_key_id, 

162 keys=_resolved_keys, 

163 grace_period_seconds=float(grace_period_seconds), 

164 ) 

165 

166 # Initialize mixin with ID generator 

167 super().__init__(ids=ids) 

168 

169 self.algorithm = algorithm 

170 self.access_expiration_hours = access_expiration_hours 

171 self.refresh_expiration_days = refresh_expiration_days 

172 self.cache_service = cache_service 

173 self.rotation_interval = rotation_interval_days * 86400 # Convert to seconds 

174 self.grace_period_seconds = grace_period_seconds 

175 

176 if logger is None: 

177 from lexigram.logging import get_logger 

178 

179 logger = get_logger(__name__) 

180 self.logger = logger.bind(manager="JWTTokenManager") 

181 

182 self._audit_logger: AuditLoggerProtocol | None = audit_logger 

183 self._binding_config = binding_config 

184 self._required_audience: str | None = required_audience 

185 self._hooks: HookRegistryProtocol | None = None 

186 

187 # Blacklist — delegates to JWTBlacklist which handles both in-process 

188 # and cache-backed revocation. 

189 self._blacklist_mgr = JWTBlacklist( 

190 cache=cache_service, 

191 cache_resolver=getattr(self, "_blacklist_cache_resolver", None), 

192 algorithm=algorithm, 

193 current_key_id_fn=lambda: self._key_store.current_key_id, 

194 access_expiration_hours=access_expiration_hours, 

195 refresh_expiration_days=refresh_expiration_days, 

196 audit_logger=audit_logger, 

197 ) 

198 

199 # Per-kid key cache: maps the JWT ``kid`` header claim to the key_id 

200 # that most recently verified a token carrying that kid. On the next 

201 # verification for the same kid the cached key_id is tried first, 

202 # avoiding redundant lookups when many rotation keys are present. 

203 # Cleared on every call to :meth:`rotate_key`. 

204 self._verified_by_key: dict[str, str] = {} 

205 

206 # Verification result cache: maps a short token hash (first 16 hex 

207 # chars of SHA-256) to the key_id that last successfully verified it. 

208 # Allows skipping the kid-based lookup entirely on repeated calls for 

209 # the same token — common in middleware that re-validates per request. 

210 # Bounded at 1 000 entries; cleared in full when the limit is hit to 

211 # avoid unbounded memory growth. Also cleared on key rotation. 

212 self._verification_cache: OrderedDict[str, str] = OrderedDict() 

213 

214 # ── Key-store properties ───────────────────────────────────────────── 

215 

216 @property 

217 def keys(self) -> dict[str, Any]: 

218 """Live view of the key material managed by the key store.""" 

219 return self._key_store.keys 

220 

221 @property 

222 def _key_meta(self) -> dict[str, Any]: 

223 """Live view of the key metadata managed by the key store.""" 

224 return self._key_store._key_meta 

225 

226 @property 

227 def current_key_id(self) -> str: 

228 """The key ID currently used for signing new tokens.""" 

229 return self._key_store.current_key_id 

230 

231 def set_blacklist_resolver(self, resolver: Any) -> None: 

232 """Attach a deferred blacklist-cache source. 

233 

234 Args: 

235 resolver: Zero-arg callable returning a ``CacheBackendProtocol`` 

236 or ``None``; invoked lazily on first revocation use. 

237 """ 

238 self._blacklist_cache_resolver = resolver 

239 self._blacklist_mgr.attach_cache_resolver(resolver) 

240 

241 def set_blacklist_cache(self, cache: Any) -> None: 

242 """Attach a cache backend to the blacklist for durable revocation. 

243 

244 Lets DI providers supply the application cache during ``boot()`` 

245 when it was not available at construction time. 

246 

247 Args: 

248 cache: Cache backend used for token/user revocation entries. 

249 """ 

250 self._blacklist_mgr.attach_cache(cache) 

251 

252 @current_key_id.setter 

253 def current_key_id(self, value: str) -> None: 

254 """Allow external callers to update the active key ID on the store.""" 

255 self._key_store.current_key_id = value 

256 

257 def __repr__(self) -> str: 

258 """Return developer-friendly string representation.""" 

259 return ( 

260 f"JWTTokenManager(algorithm={self.algorithm!r}, " 

261 f"access_expiration_hours={self.access_expiration_hours}, " 

262 f"current_key_id={self.current_key_id!r})" 

263 ) 

264 

265 def set_hook_registry(self, hooks: HookRegistryProtocol | None) -> None: 

266 """Attach an optional hook registry after provider boot wiring.""" 

267 self._hooks = hooks 

268 

269 async def rotate_key(self, new_key_id: str, new_secret: str | dict) -> None: 

270 """Rotate to a new signing key, delegating lifecycle to the key store. 

271 

272 Old keys are retained for ``grace_period_seconds`` (constructor 

273 parameter, default 3600 s) so tokens they signed remain verifiable 

274 during the overlap window. 

275 

276 Args: 

277 new_key_id: ID for new key. 

278 new_secret: New secret key (string for symmetric or dict for asymmetric). 

279 """ 

280 await self._key_store.rotate(new_key_id, new_secret) 

281 # Clear the per-kid and per-token verification caches so stale mappings 

282 # are not used after rotation (new keys may have the same kid value). 

283 self._verified_by_key.clear() 

284 self._verification_cache.clear() 

285 

286 async def _cleanup_old_keys(self) -> None: 

287 """Delegate old-key cleanup to the key store.""" 

288 await self._key_store._cleanup_old_keys() 

289 

290 def list_keys(self) -> dict[str, dict[str, Any]]: 

291 """Return current key metadata (for inspection/operations).""" 

292 return self._key_store.list_keys() 

293 

294 def _get_signing_key(self) -> str: 

295 """Return the raw signing key string for the current key ID.""" 

296 return self._key_store.get_signing_key() 

297 

298 def _get_verification_key(self, kid: str) -> str: 

299 """Return the raw verification key string for *kid*.""" 

300 return self._key_store.get_verification_key(kid) # type: ignore[return-value] 

301 

302 async def get_user_from_token( 

303 self, token: str 

304 ) -> Result[VerifiedToken, ContractsTokenError]: 

305 """Extract user information from access token. 

306 

307 Returns: 

308 ``Ok(VerifiedToken)`` if the token is a valid access token, or 

309 ``Err(TokenError)`` for expected domain failures. 

310 """ 

311 return await self.verify_token(token, "access") 

312 

313 

314__all__ = ["JWTTokenManager", "TokenBindingConfig"]