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

1import json 

2import urllib.request 

3from functools import lru_cache 

4from typing import Any, Dict, List, Optional 

5from urllib.error import URLError 

6 

7from .api_jwk import PyJWK, PyJWKSet 

8from .api_jwt import decode_complete as decode_token 

9from .exceptions import PyJWKClientConnectionError, PyJWKClientError 

10from .jwk_set_cache import JWKSetCache 

11 

12 

13class PyJWKClient: 

14 def __init__( 

15 self, 

16 uri: str, 

17 cache_keys: bool = False, 

18 max_cached_keys: int = 16, 

19 cache_jwk_set: bool = True, 

20 lifespan: int = 300, 

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

22 timeout: int = 30, 

23 ): 

24 if headers is None: 

25 headers = {} 

26 self.uri = uri 

27 self.jwk_set_cache: Optional[JWKSetCache] = None 

28 self.headers = headers 

29 self.timeout = timeout 

30 

31 if cache_jwk_set: 

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

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

34 if lifespan <= 0: 

35 raise PyJWKClientError( 

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

37 ) 

38 self.jwk_set_cache = JWKSetCache(lifespan) 

39 else: 

40 self.jwk_set_cache = None 

41 

42 if cache_keys: 

43 # Cache signing keys 

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

45 self.get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key) # type: ignore 

46 

47 def fetch_data(self) -> Any: 

48 jwk_set: Any = None 

49 try: 

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

51 with urllib.request.urlopen(r, timeout=self.timeout) as response: 

52 jwk_set = json.load(response) 

53 except (URLError, TimeoutError) as e: 

54 raise PyJWKClientConnectionError( 

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

56 ) 

57 else: 

58 return jwk_set 

59 finally: 

60 if self.jwk_set_cache is not None: 

61 self.jwk_set_cache.put(jwk_set) 

62 

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

64 data = None 

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

66 data = self.jwk_set_cache.get() 

67 

68 if data is None: 

69 data = self.fetch_data() 

70 

71 if not isinstance(data, dict): 

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

73 

74 return PyJWKSet.from_dict(data) 

75 

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

77 jwk_set = self.get_jwk_set(refresh) 

78 signing_keys = [ 

79 jwk_set_key 

80 for jwk_set_key in jwk_set.keys 

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

82 ] 

83 

84 if not signing_keys: 

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

86 

87 return signing_keys 

88 

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

90 signing_keys = self.get_signing_keys() 

91 signing_key = self.match_kid(signing_keys, kid) 

92 

93 if not signing_key: 

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

95 signing_keys = self.get_signing_keys(refresh=True) 

96 signing_key = self.match_kid(signing_keys, kid) 

97 

98 if not signing_key: 

99 raise PyJWKClientError( 

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

101 ) 

102 

103 return signing_key 

104 

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

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

107 header = unverified["header"] 

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

109 

110 @staticmethod 

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

112 signing_key = None 

113 

114 for key in signing_keys: 

115 if key.key_id == kid: 

116 signing_key = key 

117 break 

118 

119 return signing_key