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

1import json 

2import urllib.request 

3from functools import lru_cache 

4from ssl import SSLContext 

5from typing import Any, Dict, List, Optional 

6from urllib.error import URLError 

7 

8from .api_jwk import PyJWK, PyJWKSet 

9from .api_jwt import decode_complete as decode_token 

10from .exceptions import PyJWKClientConnectionError, PyJWKClientError 

11from .jwk_set_cache import JWKSetCache 

12 

13 

14class PyJWKClient: 

15 def __init__( 

16 self, 

17 uri: str, 

18 cache_keys: bool = False, 

19 max_cached_keys: int = 16, 

20 cache_jwk_set: bool = True, 

21 lifespan: int = 300, 

22 headers: Optional[Dict[str, Any]] = None, 

23 timeout: int = 30, 

24 ssl_context: Optional[SSLContext] = None, 

25 ): 

26 if headers is None: 

27 headers = {} 

28 self.uri = uri 

29 self.jwk_set_cache: Optional[JWKSetCache] = None 

30 self.headers = headers 

31 self.timeout = timeout 

32 self.ssl_context = ssl_context 

33 

34 if cache_jwk_set: 

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

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

37 if lifespan <= 0: 

38 raise PyJWKClientError( 

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

40 ) 

41 self.jwk_set_cache = JWKSetCache(lifespan) 

42 else: 

43 self.jwk_set_cache = None 

44 

45 if cache_keys: 

46 # Cache signing keys 

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

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

49 self.get_signing_key 

50 ) # type: ignore 

51 

52 def fetch_data(self) -> Any: 

53 jwk_set: Any = None 

54 try: 

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

56 with urllib.request.urlopen( 

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

58 ) as response: 

59 jwk_set = json.load(response) 

60 except (URLError, TimeoutError) as e: 

61 raise PyJWKClientConnectionError( 

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

63 ) from e 

64 else: 

65 return jwk_set 

66 finally: 

67 if self.jwk_set_cache is not None: 

68 self.jwk_set_cache.put(jwk_set) 

69 

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

71 data = None 

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

73 data = self.jwk_set_cache.get() 

74 

75 if data is None: 

76 data = self.fetch_data() 

77 

78 if not isinstance(data, dict): 

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

80 

81 return PyJWKSet.from_dict(data) 

82 

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

84 jwk_set = self.get_jwk_set(refresh) 

85 signing_keys = [ 

86 jwk_set_key 

87 for jwk_set_key in jwk_set.keys 

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

89 ] 

90 

91 if not signing_keys: 

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

93 

94 return signing_keys 

95 

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

97 signing_keys = self.get_signing_keys() 

98 signing_key = self.match_kid(signing_keys, kid) 

99 

100 if not signing_key: 

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

102 signing_keys = self.get_signing_keys(refresh=True) 

103 signing_key = self.match_kid(signing_keys, kid) 

104 

105 if not signing_key: 

106 raise PyJWKClientError( 

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

108 ) 

109 

110 return signing_key 

111 

112 def get_signing_key_from_jwt(self, token: str) -> PyJWK: 

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

114 header = unverified["header"] 

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

116 

117 @staticmethod 

118 def match_kid(signing_keys: List[PyJWK], kid: str) -> Optional[PyJWK]: 

119 signing_key = None 

120 

121 for key in signing_keys: 

122 if key.key_id == kid: 

123 signing_key = key 

124 break 

125 

126 return signing_key