D7net
Home
Console
Upload
information
Create File
Create Folder
About
Tools
:
/
proc
/
self
/
root
/
opt
/
cloudlinux
/
venv
/
lib
/
python3.11
/
site-packages
/
jwt
/
Filename :
jwks_client.py
back
Copy
import json import urllib.request from functools import lru_cache from ssl import SSLContext from typing import Any, Dict, List, Optional from urllib.error import URLError from .api_jwk import PyJWK, PyJWKSet from .api_jwt import decode_complete as decode_token from .exceptions import PyJWKClientConnectionError, PyJWKClientError from .jwk_set_cache import JWKSetCache class PyJWKClient: def __init__( self, uri: str, cache_keys: bool = False, max_cached_keys: int = 16, cache_jwk_set: bool = True, lifespan: int = 300, headers: Optional[Dict[str, Any]] = None, timeout: int = 30, ssl_context: Optional[SSLContext] = None, ): if headers is None: headers = {} self.uri = uri self.jwk_set_cache: Optional[JWKSetCache] = None self.headers = headers self.timeout = timeout self.ssl_context = ssl_context if cache_jwk_set: # Init jwt set cache with default or given lifespan. # Default lifespan is 300 seconds (5 minutes). if lifespan <= 0: raise PyJWKClientError( f'Lifespan must be greater than 0, the input is "{lifespan}"' ) self.jwk_set_cache = JWKSetCache(lifespan) else: self.jwk_set_cache = None if cache_keys: # Cache signing keys # Ignore mypy (https://github.com/python/mypy/issues/2427) self.get_signing_key = lru_cache(maxsize=max_cached_keys)(self.get_signing_key) # type: ignore def fetch_data(self) -> Any: jwk_set: Any = None try: r = urllib.request.Request(url=self.uri, headers=self.headers) with urllib.request.urlopen( r, timeout=self.timeout, context=self.ssl_context ) as response: jwk_set = json.load(response) except (URLError, TimeoutError) as e: raise PyJWKClientConnectionError( f'Fail to fetch data from the url, err: "{e}"' ) else: return jwk_set finally: if self.jwk_set_cache is not None: self.jwk_set_cache.put(jwk_set) def get_jwk_set(self, refresh: bool = False) -> PyJWKSet: data = None if self.jwk_set_cache is not None and not refresh: data = self.jwk_set_cache.get() if data is None: data = self.fetch_data() if not isinstance(data, dict): raise PyJWKClientError("The JWKS endpoint did not return a JSON object") return PyJWKSet.from_dict(data) def get_signing_keys(self, refresh: bool = False) -> List[PyJWK]: jwk_set = self.get_jwk_set(refresh) signing_keys = [ jwk_set_key for jwk_set_key in jwk_set.keys if jwk_set_key.public_key_use in ["sig", None] and jwk_set_key.key_id ] if not signing_keys: raise PyJWKClientError("The JWKS endpoint did not contain any signing keys") return signing_keys def get_signing_key(self, kid: str) -> PyJWK: signing_keys = self.get_signing_keys() signing_key = self.match_kid(signing_keys, kid) if not signing_key: # If no matching signing key from the jwk set, refresh the jwk set and try again. signing_keys = self.get_signing_keys(refresh=True) signing_key = self.match_kid(signing_keys, kid) if not signing_key: raise PyJWKClientError( f'Unable to find a signing key that matches: "{kid}"' ) return signing_key def get_signing_key_from_jwt(self, token: str) -> PyJWK: unverified = decode_token(token, options={"verify_signature": False}) header = unverified["header"] return self.get_signing_key(header.get("kid")) @staticmethod def match_kid(signing_keys: List[PyJWK], kid: str) -> Optional[PyJWK]: signing_key = None for key in signing_keys: if key.key_id == kid: signing_key = key break return signing_key