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."""
4from __future__ import annotations
6import secrets
7from typing import TYPE_CHECKING, Annotated, Any
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
24if TYPE_CHECKING:
25 from lexigram.auth.config import AuthConfig
26 from lexigram.contracts.core.di import (
27 ContainerRegistrarProtocol,
28 ContainerResolverProtocol,
29 )
31logger = get_logger(__name__)
34@inject
35class TokenProvider(Provider):
36 """JWT token management ONLY."""
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
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}
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 )
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 )
104 self.secret_key: str = resolved_secret
106 logger.info(
107 "jwt_verification_policy_boot",
108 environment=env.value,
109 mode="verified_only",
110 )
111 # ── End JWT verification policy ──────────────────────────────────────
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 )
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
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")
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 )
185 keys_to_pass = {
186 "default": {"private": private_pem, "public": public_pem},
187 }
188 current_kid = "default"
190 logger.warning(
191 "TokenProvider: using generated ephemeral RSA keypair for RS algorithm; provide persistent keys for production",
192 )
194 # Initialize token manager with keys or legacy secret
195 from pydantic import SecretStr
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 )
221 # Register with container
222 container.singleton(JWTTokenManager, lambda: self.token_manager)
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)
230 async def shutdown(self) -> None:
231 """Shutdown token provider."""
232 logger.info("TokenProvider shutdown")
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 )
248__all__ = [
249 "TokenProvider",
250 "logger",
251]