jwks_client.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. import json
  2. import urllib.request
  3. from functools import lru_cache
  4. from typing import Any, List
  5. from .api_jwk import PyJWK, PyJWKSet
  6. from .api_jwt import decode_complete as decode_token
  7. from .exceptions import PyJWKClientError
  8. class PyJWKClient:
  9. def __init__(self, uri: str, cache_keys: bool = True, max_cached_keys: int = 16):
  10. self.uri = uri
  11. if cache_keys:
  12. # Cache signing keys
  13. # Ignore mypy (https://github.com/python/mypy/issues/2427)
  14. self.get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key) # type: ignore
  15. def fetch_data(self) -> Any:
  16. with urllib.request.urlopen(self.uri) as response:
  17. return json.load(response)
  18. def get_jwk_set(self) -> PyJWKSet:
  19. data = self.fetch_data()
  20. return PyJWKSet.from_dict(data)
  21. def get_signing_keys(self) -> List[PyJWK]:
  22. jwk_set = self.get_jwk_set()
  23. signing_keys = [
  24. jwk_set_key
  25. for jwk_set_key in jwk_set.keys
  26. if jwk_set_key.public_key_use in ["sig", None] and jwk_set_key.key_id
  27. ]
  28. if not signing_keys:
  29. raise PyJWKClientError("The JWKS endpoint did not contain any signing keys")
  30. return signing_keys
  31. def get_signing_key(self, kid: str) -> PyJWK:
  32. signing_keys = self.get_signing_keys()
  33. signing_key = None
  34. for key in signing_keys:
  35. if key.key_id == kid:
  36. signing_key = key
  37. break
  38. if not signing_key:
  39. raise PyJWKClientError(
  40. f'Unable to find a signing key that matches: "{kid}"'
  41. )
  42. return signing_key
  43. def get_signing_key_from_jwt(self, token: str) -> PyJWK:
  44. unverified = decode_token(token, options={"verify_signature": False})
  45. header = unverified["header"]
  46. return self.get_signing_key(header.get("kid"))