Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/jwt/jwks_client.py: 27%

1from __future__ import annotations 

2 

3import json 

4import math 

5import threading 

6import time 

7import urllib.request 

8from functools import lru_cache 

9from ssl import SSLContext 

10from typing import Any 

11from urllib.error import HTTPError, URLError 

12from urllib.parse import urlparse 

13 

14from .api_jwk import PyJWK, PyJWKSet 

15from .api_jwt import decode_complete as decode_token 

16from .exceptions import PyJWKClientConnectionError, PyJWKClientError 

17from .jwk_set_cache import JWKSetCache 

18 

19 

20class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): 

21 def redirect_request( 

22 self, 

23 req: urllib.request.Request, 

24 fp: Any, 

25 code: int, 

26 msg: str, 

27 headers: Any, 

28 newurl: str, 

29 ) -> urllib.request.Request | None: 

30 return None 

31 

32 

33class PyJWKClient: 

34 def __init__( 

35 self, 

36 uri: str, 

37 cache_keys: bool = False, 

38 max_cached_keys: int = 16, 

39 cache_jwk_set: bool = True, 

40 lifespan: float = 300, 

41 headers: dict[str, Any] | None = None, 

42 timeout: float = 30, 

43 ssl_context: SSLContext | None = None, 

44 cooldown_duration: float = 30, 

45 ): 

46 """A client for retrieving signing keys from a JWKS endpoint. 

47 

48 ``PyJWKClient`` uses a two-tier caching system to avoid unnecessary 

49 network requests: 

50 

51 **Tier 1 — JWK Set cache** (enabled by default): 

52 Caches the entire JSON Web Key Set response from the endpoint. 

53 Controlled by: 

54 

55 - ``cache_jwk_set``: Set to ``True`` (the default) to enable this 

56 cache. When enabled, the JWK Set is fetched from the network only 

57 when the cache is empty or expired. 

58 - ``lifespan``: Time in seconds before the cached JWK Set expires. 

59 Defaults to ``300`` (5 minutes). Must be greater than 0. 

60 

61 Unknown key IDs can trigger one forced refresh after the cooldown 

62 period configured by ``cooldown_duration``. Every successful fetch 

63 starts this cooldown, including the initial fetch and cache-expiry 

64 fetches. A newly rotated key may therefore wait for the cooldown 

65 period before it is fetched; set ``cooldown_duration`` to ``0`` to 

66 disable this behavior. The cooldown is bypassed when 

67 ``cache_jwk_set`` is ``False``. 

68 

69 **Tier 2 — Signing key cache** (disabled by default): 

70 Caches individual signing keys (looked up by ``kid``) using an LRU 

71 cache with **no time-based expiration**. Keys are evicted only when 

72 the cache reaches its maximum size. Controlled by: 

73 

74 - ``cache_keys``: Set to ``True`` to enable this cache. 

75 Defaults to ``False``. 

76 - ``max_cached_keys``: Maximum number of signing keys to keep in 

77 the LRU cache. Defaults to ``16``. 

78 

79 :param uri: The URL of the JWKS endpoint. 

80 :type uri: str 

81 :param cache_keys: Enable the per-key LRU cache (Tier 2). 

82 :type cache_keys: bool 

83 :param max_cached_keys: Max entries in the signing key LRU cache. 

84 :type max_cached_keys: int 

85 :param cache_jwk_set: Enable the JWK Set response cache (Tier 1). 

86 :type cache_jwk_set: bool 

87 :param lifespan: TTL in seconds for the JWK Set cache. 

88 :type lifespan: float 

89 :param headers: Optional HTTP headers to include in requests. 

90 :type headers: dict or None 

91 :param timeout: HTTP request timeout in seconds. 

92 :type timeout: float 

93 :param ssl_context: Optional SSL context for the request. 

94 :type ssl_context: ssl.SSLContext or None 

95 :param cooldown_duration: Minimum time in seconds between forced 

96 refreshes after an unknown key ID. Defaults to ``30``. 

97 :type cooldown_duration: float 

98 """ 

99 if headers is None: 

100 headers = {} 

101 # urllib's default OpenerDirector also handles file://, ftp://, and 

102 # data: URIs. Reject anything that isn't http(s) eagerly so a caller 

103 # passing an attacker-influenced URL (e.g. taken from a `jku` token 

104 # header) can't read local files or reach other unintended schemes. 

105 scheme = urlparse(uri).scheme.lower() 

106 if scheme not in ("http", "https"): 

107 raise PyJWKClientError( 

108 f"Invalid JWKS URI scheme {scheme!r}: only 'http' and 'https' " 

109 f"are supported." 

110 ) 

111 self.uri = uri 

112 self.jwk_set_cache: JWKSetCache | None = None 

113 self.headers = headers 

114 self.timeout = timeout 

115 self.ssl_context = ssl_context 

116 if cooldown_duration < 0: 

117 raise PyJWKClientError( 

118 "Cooldown duration must be greater than or equal to 0, " 

119 f'the input is "{cooldown_duration}"' 

120 ) 

121 if not math.isfinite(cooldown_duration): 

122 raise PyJWKClientError( 

123 f'Cooldown duration must be finite, the input is "{cooldown_duration}"' 

124 ) 

125 self.cooldown_duration = cooldown_duration 

126 self._last_successful_fetch: float | None = None 

127 self._client_lock = threading.RLock() 

128 

129 if cache_jwk_set: 

130 # Init jwt set cache with default or given lifespan. 

131 # Default lifespan is 300 seconds (5 minutes). 

132 if lifespan <= 0: 

133 raise PyJWKClientError( 

134 f'Lifespan must be greater than 0, the input is "{lifespan}"' 

135 ) 

136 self.jwk_set_cache = JWKSetCache(lifespan) 

137 else: 

138 self.jwk_set_cache = None 

139 

140 if cache_keys: 

141 # Cache signing keys 

142 get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key) 

143 # Ignore mypy (https://github.com/python/mypy/issues/2427) 

144 self.get_signing_key = get_signing_key # type: ignore[method-assign] 

145 

146 def fetch_data(self) -> Any: 

147 """Fetch the JWK Set from the JWKS endpoint. 

148 

149 Makes an HTTP request to the configured ``uri`` and returns the 

150 parsed JSON response. If the JWK Set cache is enabled, the 

151 response is stored in the cache. 

152 

153 :returns: The parsed JWK Set as a dictionary. 

154 :raises PyJWKClientConnectionError: If the HTTP request fails. 

155 """ 

156 try: 

157 r = urllib.request.Request(url=self.uri, headers=self.headers) 

158 handlers: list[Any] = [_NoRedirectHandler()] 

159 if self.ssl_context is not None: 

160 handlers.append(urllib.request.HTTPSHandler(context=self.ssl_context)) 

161 opener = urllib.request.build_opener(*handlers) 

162 with opener.open(r, timeout=self.timeout) as response: 

163 jwk_set = json.load(response) 

164 except (URLError, TimeoutError) as e: 

165 if isinstance(e, HTTPError): 

166 e.close() 

167 raise PyJWKClientConnectionError( 

168 f'Fail to fetch data from the url, err: "{e}"' 

169 ) from e 

170 

171 # Only update the cache on a successful fetch. Writing in a 

172 # `finally` block with `jwk_set=None` on error clears any 

173 # previously-cached JWKS, turning a transient outage into a cache 

174 # wipe that breaks legitimate auth. 

175 if self.jwk_set_cache is not None: 

176 self.jwk_set_cache.put(jwk_set) 

177 self._last_successful_fetch = time.monotonic() 

178 return jwk_set 

179 

180 def get_jwk_set(self, refresh: bool = False) -> PyJWKSet: 

181 """Return the JWK Set, using the cache when available. 

182 

183 :param refresh: Force a fresh fetch from the endpoint, bypassing 

184 the cache. 

185 :type refresh: bool 

186 :returns: The JWK Set. 

187 :rtype: PyJWKSet 

188 :raises PyJWKClientError: If the endpoint does not return a JSON 

189 object. 

190 """ 

191 data = None 

192 if self.jwk_set_cache is not None and not refresh: 

193 data = self.jwk_set_cache.get() 

194 

195 if data is None: 

196 data = self.fetch_data() 

197 

198 if not isinstance(data, dict): 

199 raise PyJWKClientError("The JWKS endpoint did not return a JSON object") 

200 

201 return PyJWKSet.from_dict(data) 

202 

203 def get_signing_keys(self, refresh: bool = False) -> list[PyJWK]: 

204 """Return all signing keys from the JWK Set. 

205 

206 Filters the JWK Set to keys whose ``use`` is ``"sig"`` (or 

207 unspecified) and that have a ``kid``. 

208 

209 :param refresh: Force a fresh fetch from the endpoint, bypassing 

210 the cache. 

211 :type refresh: bool 

212 :returns: A list of signing keys. 

213 :rtype: list[PyJWK] 

214 :raises PyJWKClientError: If no signing keys are found. 

215 """ 

216 jwk_set = self.get_jwk_set(refresh) 

217 return self._get_signing_keys_from_jwk_set(jwk_set) 

218 

219 @staticmethod 

220 def _get_signing_keys_from_jwk_set(jwk_set: PyJWKSet) -> list[PyJWK]: 

221 signing_keys = [ 

222 jwk_set_key 

223 for jwk_set_key in jwk_set.keys 

224 if jwk_set_key.public_key_use in ["sig", None] and jwk_set_key.key_id 

225 ] 

226 

227 if not signing_keys: 

228 raise PyJWKClientError("The JWKS endpoint did not contain any signing keys") 

229 

230 return signing_keys 

231 

232 def get_signing_key(self, kid: str) -> PyJWK: 

233 """Return the signing key matching the given ``kid``. 

234 

235 If no match is found in the current JWK Set, the set is refreshed 

236 from the endpoint and the lookup is retried once when the refresh 

237 cooldown permits it. 

238 

239 :param kid: The key ID to look up. 

240 :type kid: str 

241 :returns: The matching signing key. 

242 :rtype: PyJWK 

243 :raises PyJWKClientError: If no matching key is found after 

244 refreshing. 

245 """ 

246 with self._client_lock: 

247 signing_keys = self.get_signing_keys() 

248 signing_key = self.match_kid(signing_keys, kid) 

249 

250 if not signing_key: 

251 cooling_down = ( 

252 self.jwk_set_cache is not None 

253 and self._last_successful_fetch is not None 

254 and time.monotonic() - self._last_successful_fetch 

255 < self.cooldown_duration 

256 ) 

257 if not cooling_down: 

258 signing_keys = self.get_signing_keys(refresh=True) 

259 self._last_successful_fetch = time.monotonic() 

260 signing_key = self.match_kid(signing_keys, kid) 

261 

262 if not signing_key: 

263 raise PyJWKClientError( 

264 f'Unable to find a signing key that matches: "{kid}"' 

265 ) 

266 

267 return signing_key 

268 

269 def get_signing_key_from_jwt(self, token: str | bytes) -> PyJWK: 

270 """Return the signing key for a JWT by reading its ``kid`` header. 

271 

272 Extracts the ``kid`` from the token's unverified header and 

273 delegates to :meth:`get_signing_key`. 

274 

275 :param token: The encoded JWT. 

276 :type token: str or bytes 

277 :returns: The matching signing key. 

278 :rtype: PyJWK 

279 """ 

280 unverified = decode_token(token, options={"verify_signature": False}) 

281 header = unverified["header"] 

282 return self.get_signing_key(header.get("kid")) 

283 

284 @staticmethod 

285 def match_kid(signing_keys: list[PyJWK], kid: str) -> PyJWK | None: 

286 """Find a key in *signing_keys* that matches *kid*. 

287 

288 :param signing_keys: The list of keys to search. 

289 :type signing_keys: list[PyJWK] 

290 :param kid: The key ID to match. 

291 :type kid: str 

292 :returns: The matching key, or ``None`` if not found. 

293 :rtype: PyJWK or None 

294 """ 

295 signing_key = None 

296 

297 for key in signing_keys: 

298 if key.key_id == kid: 

299 signing_key = key 

300 break 

301 

302 return signing_key