From 4058ae0d03522e8fa751da0b49018842a7bf7e91 Mon Sep 17 00:00:00 2001 From: Stephen Rosen Date: Wed, 29 Jun 2022 14:48:25 +0000 Subject: [PATCH 1/2] Expose get_algorithm_by_name as new method Looking up an algorithm by name is used internally for signature generation. This encapsulates that functionality in a dedicated method and adds it to the public API. No new tests are needed to exercise the functionality. Rationale: 1. Inside of PyJWS, this improves the code. The KeyError handler is better scoped and the signing code reads more directly. 2. This is part of the path to supporting OIDC at_hash validation as a use-case (see: #295, #296, #314). This is arguably sufficient to consider that use-case supported and close it. However, it is an improvement and step in the right direction in either case. A minor change was needed to satisfy mypy, as a union-typed variable does not narrow its type based on assignments. The easiest resolution is to use a new name, in this case, simply `algorithm -> algorithm_`. --- CHANGELOG.rst | 2 ++ jwt/__init__.py | 2 ++ jwt/api_jws.py | 40 +++++++++++++++++++++++++--------------- 3 files changed, 29 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c11ce06a..81d421e2 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,8 @@ Fixed Added ~~~~~ - Add to_jwk static method to ECAlgorithm by @leonsmith in https://github.com/jpadilla/pyjwt/pull/732 +- Add ``get_algorithm_by_name`` as a method of ``PyJWS`` objects, and expose + the global PyJWS method as part of the public API `v2.4.0 `__ ----------------------------------------------------------------------- diff --git a/jwt/__init__.py b/jwt/__init__.py index 6b3f8ab1..a96cc6ee 100644 --- a/jwt/__init__.py +++ b/jwt/__init__.py @@ -1,6 +1,7 @@ from .api_jwk import PyJWK, PyJWKSet from .api_jws import ( PyJWS, + get_algorithm_by_name, get_unverified_header, register_algorithm, unregister_algorithm, @@ -51,6 +52,7 @@ "get_unverified_header", "register_algorithm", "unregister_algorithm", + "get_algorithm_by_name", # Exceptions "DecodeError", "ExpiredSignatureError", diff --git a/jwt/api_jws.py b/jwt/api_jws.py index 16ec846c..833d46aa 100644 --- a/jwt/api_jws.py +++ b/jwt/api_jws.py @@ -73,6 +73,23 @@ def get_algorithms(self): """ return list(self._valid_algs) + def get_algorithm_by_name(self, alg_name: str) -> Algorithm: + """ + For a given string name, return the matching Algorithm object. + + Example usage: + + >>> jws_obj.get_algorithm_by_name("RS256") + """ + try: + return self._algorithms[alg_name] + except KeyError as e: + if not has_crypto and alg_name in requires_cryptography: + raise NotImplementedError( + f"Algorithm '{alg_name}' could not be found. Do you have cryptography installed?" + ) from e + raise NotImplementedError("Algorithm not supported") from e + def encode( self, payload: bytes, @@ -84,21 +101,21 @@ def encode( ) -> str: segments = [] - if algorithm is None: - algorithm = "none" + # declare a new var to narrow the type for type checkers + algorithm_: str = algorithm if algorithm is not None else "none" # Prefer headers values if present to function parameters. if headers: headers_alg = headers.get("alg") if headers_alg: - algorithm = 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] + header = {"typ": self.header_typ, "alg": algorithm_} # type: Dict[str, Any] if headers: self._validate_headers(headers) @@ -128,17 +145,9 @@ def encode( # 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 + alg_obj = self.get_algorithm_by_name(algorithm_) + key = alg_obj.prepare_key(key) + signature = alg_obj.sign(signing_input, key) segments.append(base64url_encode(signature)) @@ -286,4 +295,5 @@ def _validate_kid(self, kid): decode = _jws_global_obj.decode register_algorithm = _jws_global_obj.register_algorithm unregister_algorithm = _jws_global_obj.unregister_algorithm +get_algorithm_by_name = _jws_global_obj.get_algorithm_by_name get_unverified_header = _jws_global_obj.get_unverified_header From e612fbd19006b068afc29d646e282b80d0a909d2 Mon Sep 17 00:00:00 2001 From: Stephen Rosen Date: Wed, 29 Jun 2022 15:27:19 +0000 Subject: [PATCH 2/2] Use get_algorithm_by_name in _verify_signature Rather than catching the KeyError from a dict lookup, catch the NotImplementedError raised by get_algorithm_by_name. This changes the exception seen in the cause under exception chaining but otherwise has no public-facing impact. --- jwt/api_jws.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/jwt/api_jws.py b/jwt/api_jws.py index 833d46aa..f8c8048d 100644 --- a/jwt/api_jws.py +++ b/jwt/api_jws.py @@ -271,14 +271,13 @@ def _verify_signature( 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: + alg_obj = self.get_algorithm_by_name(alg) + except NotImplementedError as e: raise InvalidAlgorithmError("Algorithm not supported") from e + key = alg_obj.prepare_key(key) + + if not alg_obj.verify(signing_input, key, signature): + raise InvalidSignatureError("Signature verification failed") def _validate_headers(self, headers): if "kid" in headers: