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 urllib.request 

5from functools import lru_cache 

6from ssl import SSLContext 

7from typing import Any 

8from urllib.error import URLError 

9 

10from .api_jwk import PyJWK, PyJWKSet 

11from .api_jwt import decode_complete as decode_token 

12from .exceptions import PyJWKClientConnectionError, PyJWKClientError 

13from .jwk_set_cache import JWKSetCache 

14 

15 

16class PyJWKClient: 

17 def __init__( 

18 self, 

19 uri: str, 

20 cache_keys: bool = False, 

21 max_cached_keys: int = 16, 

22 cache_jwk_set: bool = True, 

23 lifespan: float = 300, 

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

25 timeout: float = 30, 

26 ssl_context: SSLContext | None = None, 

27 ): 

28 if headers is None: 

29 headers = {} 

30 self.uri = uri 

31 self.jwk_set_cache: JWKSetCache | None = None 

32 self.headers = headers 

33 self.timeout = timeout 

34 self.ssl_context = ssl_context 

35 

36 if cache_jwk_set: 

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

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

39 if lifespan <= 0: 

40 raise PyJWKClientError( 

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

42 ) 

43 self.jwk_set_cache = JWKSetCache(lifespan) 

44 else: 

45 self.jwk_set_cache = None 

46 

47 if cache_keys: 

48 # Cache signing keys 

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

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

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

52 

53 def fetch_data(self) -> Any: 

54 jwk_set: Any = None 

55 try: 

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

57 with urllib.request.urlopen( 

58 r, timeout=self.timeout, context=self.ssl_context 

59 ) as response: 

60 jwk_set = json.load(response) 

61 except (URLError, TimeoutError) as e: 

62 raise PyJWKClientConnectionError( 

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

64 ) from e 

65 else: 

66 return jwk_set 

67 finally: 

68 if self.jwk_set_cache is not None: 

69 self.jwk_set_cache.put(jwk_set) 

70 

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

72 data = None 

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

74 data = self.jwk_set_cache.get() 

75 

76 if data is None: 

77 data = self.fetch_data() 

78 

79 if not isinstance(data, dict): 

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

81 

82 return PyJWKSet.from_dict(data) 

83 

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

85 jwk_set = self.get_jwk_set(refresh) 

86 signing_keys = [ 

87 jwk_set_key 

88 for jwk_set_key in jwk_set.keys 

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

90 ] 

91 

92 if not signing_keys: 

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

94 

95 return signing_keys 

96 

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

98 signing_keys = self.get_signing_keys() 

99 signing_key = self.match_kid(signing_keys, kid) 

100 

101 if not signing_key: 

102 # If no matching signing key from the jwk set, refresh the jwk set and try again. 

103 signing_keys = self.get_signing_keys(refresh=True) 

104 signing_key = self.match_kid(signing_keys, kid) 

105 

106 if not signing_key: 

107 raise PyJWKClientError( 

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

109 ) 

110 

111 return signing_key 

112 

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

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

115 header = unverified["header"] 

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

117 

118 @staticmethod 

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

120 signing_key = None 

121 

122 for key in signing_keys: 

123 if key.key_id == kid: 

124 signing_key = key 

125 break 

126 

127 return signing_key