Coverage for src/lexigram/auth/di/sub_providers/token_provider.py: 76%

72 statements  

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

1# lexigram/auth/providers/token_provider.py 

2"""Token management provider - handles JWT tokens only.""" 

3 

4from __future__ import annotations 

5 

6import secrets 

7from typing import TYPE_CHECKING, Annotated, Any 

8 

9from lexigram.auth import constants as const 

10from lexigram.auth.authn.jwt import JWTTokenManager 

11from lexigram.contracts.core import ( 

12 HealthCheckResult, 

13 HealthStatus, 

14 HookRegistryProtocol, 

15 ProviderPriority, 

16) 

17from lexigram.contracts.core.config import Environment 

18from lexigram.contracts.exceptions import ConfigurationError 

19from lexigram.di.decorators import inject 

20from lexigram.di.markers import Inject 

21from lexigram.di.provider import Provider 

22from lexigram.logging import get_logger 

23 

24if TYPE_CHECKING: 

25 from lexigram.auth.config import AuthConfig 

26 from lexigram.contracts.core.di import ( 

27 ContainerRegistrarProtocol, 

28 ContainerResolverProtocol, 

29 ) 

30 

31logger = get_logger(__name__) 

32 

33 

34@inject 

35class TokenProvider(Provider): 

36 """JWT token management ONLY.""" 

37 

38 def __init__( 

39 self, 

40 config: Annotated[AuthConfig, Inject] | None = None, 

41 secret_key: str | None = None, 

42 jwt_algorithm: str | None = None, 

43 jwt_access_expiration_hours: int | None = None, 

44 jwt_refresh_expiration_days: int | None = None, 

45 **kwargs: Any, 

46 ) -> None: 

47 super().__init__(name="tokens", priority=ProviderPriority.SECURITY) 

48 token_config = config.token if config else None 

49 

50 # ── JWT verification policy ────────────────────────────────────────── 

51 # Verified-only decoding is enforced everywhere. A missing secret is 

52 # fatal in PRODUCTION/STAGING; in DEVELOPMENT an ephemeral secret is 

53 # generated so signature verification never needs to be disabled. 

54 env = Environment.from_env() 

55 _STRICT_ENVS = {Environment.PRODUCTION, Environment.STAGING} 

56 

57 # Secret resolution: 

58 # 1. Explicit secret_key argument → use it directly (verified mode). 

59 # 2. Config provided with a secret → use it (verified mode). 

60 # 3. Config provided but secret is absent → policy enforcement below. 

61 # 4. No config at all (config=None) → generate ephemeral secret (verified mode, 

62 # zero-config dev/test usage documented in AuthModule.configure(None)). 

63 if secret_key: 

64 resolved_secret: str | None = secret_key 

65 elif token_config is not None: 

66 resolved_secret = ( 

67 token_config.secret_key.get_secret_value() 

68 if token_config.secret_key 

69 else None 

70 ) 

71 else: 

72 # No config at all: ephemeral key for zero-config dev/test. 

73 if env in _STRICT_ENVS: 

74 raise ConfigurationError( 

75 f"CRITICAL SECURITY: JWT secret_key is required in {env.value.upper()} " 

76 "but no AuthConfig was provided. " 

77 "Set LEX_AUTH__TOKEN__SECRET_KEY (or pass a configured AuthConfig)." 

78 ) 

79 resolved_secret = secrets.token_urlsafe(32) 

80 logger.warning( 

81 "jwt_ephemeral_secret_generated", 

82 environment=env.value, 

83 reason="No AuthConfig supplied; using a generated ephemeral JWT secret. " 

84 "Tokens will be invalidated on restart. Provide a stable secret for production.", 

85 ) 

86 

87 if resolved_secret is None: 

88 # Config was provided but secret is absent — apply policy. 

89 if env in _STRICT_ENVS: 

90 raise ConfigurationError( 

91 f"CRITICAL SECURITY: JWT secret_key is required in {env.value.upper()} " 

92 "but none was provided. " 

93 "Set LEX_AUTH__TOKEN__SECRET_KEY (or token.secret_key in config)." 

94 ) 

95 resolved_secret = secrets.token_urlsafe(32) 

96 logger.warning( 

97 "jwt_ephemeral_secret_generated", 

98 environment=env.value, 

99 reason="No JWT secret configured; using a generated ephemeral JWT secret. " 

100 "Signature verification stays enabled. Tokens are invalidated on " 

101 "restart; set LEX_AUTH__TOKEN__SECRET_KEY for stable dev secrets.", 

102 ) 

103 

104 self.secret_key: str = resolved_secret 

105 

106 logger.info( 

107 "jwt_verification_policy_boot", 

108 environment=env.value, 

109 mode="verified_only", 

110 ) 

111 # ── End JWT verification policy ────────────────────────────────────── 

112 

113 self.jwt_algorithm = jwt_algorithm or ( 

114 token_config.algorithm if token_config else const.DEFAULT_TOKEN_ALGORITHM 

115 ) 

116 self.jwt_access_expiration_hours: int = int( 

117 jwt_access_expiration_hours 

118 or ( # type: ignore[arg-type] 

119 getattr(token_config, "access_expiration_hours", 1) 

120 if token_config 

121 else 1 

122 ) 

123 ) 

124 self.jwt_refresh_expiration_days: int = int( 

125 jwt_refresh_expiration_days 

126 or ( # type: ignore[arg-type] 

127 getattr( 

128 token_config, 

129 "refresh_expiration_days", 

130 const.DEFAULT_REFRESH_TOKEN_EXPIRE_DAYS, 

131 ) 

132 if token_config 

133 else const.DEFAULT_REFRESH_TOKEN_EXPIRE_DAYS 

134 ) 

135 ) 

136 self.jwt_key_rotation_grace_period_seconds: int = int( 

137 getattr( 

138 token_config, 

139 "key_rotation_grace_period_seconds", 

140 const.DEFAULT_JWT_KEY_ROTATION_GRACE_PERIOD_SECONDS, 

141 ) 

142 if token_config 

143 else const.DEFAULT_JWT_KEY_ROTATION_GRACE_PERIOD_SECONDS 

144 ) 

145 

146 async def register(self, container: ContainerRegistrarProtocol) -> None: 

147 """Register token services with the container.""" 

148 # If RSA algorithm chosen and secret is not an RSA PEM, generate an ephemeral keypair 

149 keys_to_pass = None 

150 current_kid = None 

151 if self.jwt_algorithm.startswith("RS"): 

152 # If secret_key looks like a PEM private key, use it 

153 if self.secret_key and "-----BEGIN" in self.secret_key: 

154 keys_to_pass = {"default": {"private": self.secret_key}} 

155 current_kid = "default" 

156 else: 

157 # Generate ephemeral RSA keypair 

158 try: 

159 from cryptography.hazmat.primitives import serialization 

160 from cryptography.hazmat.primitives.asymmetric import rsa 

161 except (ImportError, ModuleNotFoundError): 

162 raise RuntimeError( 

163 "cryptography is required for RS algorithms (install cryptography)", 

164 ) from None 

165 

166 private_key = rsa.generate_private_key( 

167 public_exponent=65537, 

168 key_size=2048, 

169 ) 

170 private_pem = private_key.private_bytes( 

171 encoding=serialization.Encoding.PEM, 

172 format=serialization.PrivateFormat.PKCS8, 

173 encryption_algorithm=serialization.NoEncryption(), 

174 ).decode("utf-8") 

175 

176 public_pem = ( 

177 private_key.public_key() 

178 .public_bytes( 

179 encoding=serialization.Encoding.PEM, 

180 format=serialization.PublicFormat.SubjectPublicKeyInfo, 

181 ) 

182 .decode("utf-8") 

183 ) 

184 

185 keys_to_pass = { 

186 "default": {"private": private_pem, "public": public_pem}, 

187 } 

188 current_kid = "default" 

189 

190 logger.warning( 

191 "TokenProvider: using generated ephemeral RSA keypair for RS algorithm; provide persistent keys for production", 

192 ) 

193 

194 # Initialize token manager with keys or legacy secret 

195 from pydantic import SecretStr 

196 

197 if keys_to_pass is not None and current_kid is not None: 

198 self.token_manager = JWTTokenManager( 

199 current_key_id=current_kid, 

200 keys=keys_to_pass, # type: ignore[arg-type] 

201 algorithm=self.jwt_algorithm, 

202 access_expiration_hours=self.jwt_access_expiration_hours, 

203 refresh_expiration_days=self.jwt_refresh_expiration_days, 

204 grace_period_seconds=self.jwt_key_rotation_grace_period_seconds, 

205 ) 

206 else: 

207 current_key = ( 

208 SecretStr(self.secret_key) 

209 if isinstance(self.secret_key, str) 

210 else self.secret_key 

211 ) 

212 self.token_manager = JWTTokenManager( 

213 current_key_id="default", 

214 keys={"default": current_key}, 

215 algorithm=self.jwt_algorithm, 

216 access_expiration_hours=self.jwt_access_expiration_hours, 

217 refresh_expiration_days=self.jwt_refresh_expiration_days, 

218 grace_period_seconds=self.jwt_key_rotation_grace_period_seconds, 

219 ) 

220 

221 # Register with container 

222 container.singleton(JWTTokenManager, lambda: self.token_manager) 

223 

224 async def boot(self, container: ContainerResolverProtocol) -> None: 

225 """Initialize token provider.""" 

226 logger.info("TokenProvider started") 

227 hooks = await container.resolve_optional(HookRegistryProtocol) 

228 self.token_manager.set_hook_registry(hooks) 

229 

230 async def shutdown(self) -> None: 

231 """Shutdown token provider.""" 

232 logger.info("TokenProvider shutdown") 

233 

234 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult: 

235 """Check token provider health.""" 

236 return HealthCheckResult( 

237 component=self.name, 

238 status=HealthStatus.HEALTHY, 

239 details={ 

240 "service": "tokens", 

241 "algorithm": self.jwt_algorithm, 

242 "access_expiration_hours": self.jwt_access_expiration_hours, 

243 "refresh_expiration_days": self.jwt_refresh_expiration_days, 

244 }, 

245 ) 

246 

247 

248__all__ = [ 

249 "TokenProvider", 

250 "logger", 

251]