123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289 |
- import binascii
- import json
- from collections.abc import Mapping
- from typing import Any, Dict, List, Optional, Type
- from .algorithms import (
- Algorithm,
- get_default_algorithms,
- has_crypto,
- requires_cryptography,
- )
- from .exceptions import (
- DecodeError,
- InvalidAlgorithmError,
- InvalidSignatureError,
- InvalidTokenError,
- )
- from .utils import base64url_decode, base64url_encode
- class PyJWS:
- header_typ = "JWT"
- def __init__(self, algorithms=None, options=None):
- self._algorithms = get_default_algorithms()
- self._valid_algs = (
- set(algorithms) if algorithms is not None else set(self._algorithms)
- )
- # Remove algorithms that aren't on the whitelist
- for key in list(self._algorithms.keys()):
- if key not in self._valid_algs:
- del self._algorithms[key]
- if options is None:
- options = {}
- self.options = {**self._get_default_options(), **options}
- @staticmethod
- def _get_default_options():
- return {"verify_signature": True}
- def register_algorithm(self, alg_id, alg_obj):
- """
- Registers a new Algorithm for use when creating and verifying tokens.
- """
- if alg_id in self._algorithms:
- raise ValueError("Algorithm already has a handler.")
- if not isinstance(alg_obj, Algorithm):
- raise TypeError("Object is not of type `Algorithm`")
- self._algorithms[alg_id] = alg_obj
- self._valid_algs.add(alg_id)
- def unregister_algorithm(self, alg_id):
- """
- Unregisters an Algorithm for use when creating and verifying tokens
- Throws KeyError if algorithm is not registered.
- """
- if alg_id not in self._algorithms:
- raise KeyError(
- "The specified algorithm could not be removed"
- " because it is not registered."
- )
- del self._algorithms[alg_id]
- self._valid_algs.remove(alg_id)
- def get_algorithms(self):
- """
- Returns a list of supported values for the 'alg' parameter.
- """
- return list(self._valid_algs)
- def encode(
- self,
- payload: bytes,
- key: str,
- algorithm: Optional[str] = "HS256",
- headers: Optional[Dict] = None,
- json_encoder: Optional[Type[json.JSONEncoder]] = None,
- is_payload_detached: bool = False,
- ) -> str:
- segments = []
- if algorithm is None:
- algorithm = "none"
- # Prefer headers values if present to function parameters.
- if headers:
- headers_alg = headers.get("alg")
- if headers_alg:
- algorithm = headers["alg"]
- headers_b64 = headers.get("b64")
- if headers_b64 is False:
- is_payload_detached = True
- # Header
- header = {"typ": self.header_typ, "alg": algorithm} # type: Dict[str, Any]
- if headers:
- self._validate_headers(headers)
- header.update(headers)
- if not header["typ"]:
- del header["typ"]
- if is_payload_detached:
- header["b64"] = False
- elif "b64" in header:
- # True is the standard value for b64, so no need for it
- del header["b64"]
- json_header = json.dumps(
- header, separators=(",", ":"), cls=json_encoder
- ).encode()
- segments.append(base64url_encode(json_header))
- if is_payload_detached:
- msg_payload = payload
- else:
- msg_payload = base64url_encode(payload)
- segments.append(msg_payload)
- # Segments
- signing_input = b".".join(segments)
- try:
- alg_obj = self._algorithms[algorithm]
- key = alg_obj.prepare_key(key)
- signature = alg_obj.sign(signing_input, key)
- except KeyError as e:
- if not has_crypto and algorithm in requires_cryptography:
- raise NotImplementedError(
- f"Algorithm '{algorithm}' could not be found. Do you have cryptography installed?"
- ) from e
- raise NotImplementedError("Algorithm not supported") from e
- segments.append(base64url_encode(signature))
- # Don't put the payload content inside the encoded token when detached
- if is_payload_detached:
- segments[1] = b""
- encoded_string = b".".join(segments)
- return encoded_string.decode("utf-8")
- def decode_complete(
- self,
- jwt: str,
- key: str = "",
- algorithms: Optional[List[str]] = None,
- options: Optional[Dict] = None,
- detached_payload: Optional[bytes] = None,
- **kwargs,
- ) -> Dict[str, Any]:
- if options is None:
- options = {}
- merged_options = {**self.options, **options}
- verify_signature = merged_options["verify_signature"]
- if verify_signature and not algorithms:
- raise DecodeError(
- 'It is required that you pass in a value for the "algorithms" argument when calling decode().'
- )
- payload, signing_input, header, signature = self._load(jwt)
- if header.get("b64", True) is False:
- if detached_payload is None:
- raise DecodeError(
- 'It is required that you pass in a value for the "detached_payload" argument to decode a message having the b64 header set to false.'
- )
- payload = detached_payload
- signing_input = b".".join([signing_input.rsplit(b".", 1)[0], payload])
- if verify_signature:
- self._verify_signature(signing_input, header, signature, key, algorithms)
- return {
- "payload": payload,
- "header": header,
- "signature": signature,
- }
- def decode(
- self,
- jwt: str,
- key: str = "",
- algorithms: Optional[List[str]] = None,
- options: Optional[Dict] = None,
- **kwargs,
- ) -> str:
- decoded = self.decode_complete(jwt, key, algorithms, options, **kwargs)
- return decoded["payload"]
- def get_unverified_header(self, jwt):
- """Returns back the JWT header parameters as a dict()
- Note: The signature is not verified so the header parameters
- should not be fully trusted until signature verification is complete
- """
- headers = self._load(jwt)[2]
- self._validate_headers(headers)
- return headers
- def _load(self, jwt):
- if isinstance(jwt, str):
- jwt = jwt.encode("utf-8")
- if not isinstance(jwt, bytes):
- raise DecodeError(f"Invalid token type. Token must be a {bytes}")
- try:
- signing_input, crypto_segment = jwt.rsplit(b".", 1)
- header_segment, payload_segment = signing_input.split(b".", 1)
- except ValueError as err:
- raise DecodeError("Not enough segments") from err
- try:
- header_data = base64url_decode(header_segment)
- except (TypeError, binascii.Error) as err:
- raise DecodeError("Invalid header padding") from err
- try:
- header = json.loads(header_data)
- except ValueError as e:
- raise DecodeError(f"Invalid header string: {e}") from e
- if not isinstance(header, Mapping):
- raise DecodeError("Invalid header string: must be a json object")
- try:
- payload = base64url_decode(payload_segment)
- except (TypeError, binascii.Error) as err:
- raise DecodeError("Invalid payload padding") from err
- try:
- signature = base64url_decode(crypto_segment)
- except (TypeError, binascii.Error) as err:
- raise DecodeError("Invalid crypto padding") from err
- return (payload, signing_input, header, signature)
- def _verify_signature(
- self,
- signing_input,
- header,
- signature,
- key="",
- algorithms=None,
- ):
- alg = header.get("alg")
- if algorithms is not None and alg not in algorithms:
- raise InvalidAlgorithmError("The specified alg value is not allowed")
- try:
- alg_obj = self._algorithms[alg]
- key = alg_obj.prepare_key(key)
- if not alg_obj.verify(signing_input, key, signature):
- raise InvalidSignatureError("Signature verification failed")
- except KeyError as e:
- raise InvalidAlgorithmError("Algorithm not supported") from e
- def _validate_headers(self, headers):
- if "kid" in headers:
- self._validate_kid(headers["kid"])
- def _validate_kid(self, kid):
- if not isinstance(kid, str):
- raise InvalidTokenError("Key ID header parameter must be a string")
- _jws_global_obj = PyJWS()
- encode = _jws_global_obj.encode
- decode_complete = _jws_global_obj.decode_complete
- decode = _jws_global_obj.decode
- register_algorithm = _jws_global_obj.register_algorithm
- unregister_algorithm = _jws_global_obj.unregister_algorithm
- get_unverified_header = _jws_global_obj.get_unverified_header
|