Split PyJWT/PyJWS classes to tighten type interfaces (jpadilla#559)

The class PyJWT was previously a subclass of PyJWS. However, this
combination does not follow the Liskov substitution principle. That is,
using PyJWT in place of a PyJWS would not produce correct results or
follow type contracts.

While these classes look to share a common interface it doesn't go
beyond the method names "encode" and "decode" and so is merely
superficial.

The classes have been split into two. PyJWT now uses composition instead
of inheritance to achieve the desired behavior. Splitting the classes in
this way allowed for precising the type interfaces.

The complete parameter to .decode() has been removed. This argument was
used to alter the return type of .decode(). Now, there are two different
methods with more explicit return types and values. The new method name
is .decode_complete(). This fills the previous role filled by
.decode(..., complete=True).

Closes jpadilla#554, jpadilla#396, jpadilla#394

Co-authored-by: Sam Bull <git@sambull.org>

Co-authored-by: Sam Bull <git@sambull.org>

Author: jdufresne

jwt/__init__.py

@@ -1,12 +1,10 @@
-from .api_jws import PyJWS
-from .api_jwt import (
-    PyJWT,
-    decode,
-    encode,
+from .api_jws import (
+    PyJWS,
     get_unverified_header,
     register_algorithm,
     unregister_algorithm,
 )
+from .api_jwt import PyJWT, decode, encode
 from .exceptions import (
     DecodeError,
     ExpiredSignatureError,

jwt/api_jws.py

@@ -1,7 +1,7 @@
 import binascii
 import json
 from collections.abc import Mapping
-from typing import Dict, List, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type
 
 from .algorithms import (
     Algorithm,
@@ -77,7 +77,7 @@ def get_algorithms(self):
 
     def encode(
         self,
-        payload: Union[Dict, bytes],
+        payload: bytes,
         key: str,
         algorithm: str = "HS256",
         headers: Optional[Dict] = None,
@@ -127,15 +127,14 @@ def encode(
 
         return encoded_string.decode("utf-8")
 
-    def decode(
+    def decode_complete(
         self,
         jwt: str,
         key: str = "",
         algorithms: List[str] = None,
         options: Dict = None,
-        complete: bool = False,
         **kwargs,
-    ):
+    ) -> Dict[str, Any]:
         if options is None:
             options = {}
         merged_options = {**self.options, **options}
@@ -153,14 +152,22 @@ def decode(
                 payload, signing_input, header, signature, key, algorithms
             )
 
-        if complete:
-            return {
-                "payload": payload,
-                "header": header,
-                "signature": signature,
-            }
+        return {
+            "payload": payload,
+            "header": header,
+            "signature": signature,
+        }
 
-        return payload
+    def decode(
+        self,
+        jwt: str,
+        key: str = "",
+        algorithms: List[str] = None,
+        options: 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()
@@ -249,6 +256,7 @@ def _validate_kid(self, kid):
 
 _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

jwt/api_jwt.py

@@ -4,7 +4,7 @@
 from datetime import datetime, timedelta
 from typing import Any, Dict, List, Optional, Type, Union
 
-from .api_jws import PyJWS
+from . import api_jws
 from .exceptions import (
     DecodeError,
     ExpiredSignatureError,
@@ -16,8 +16,11 @@
 )
 
 
-class PyJWT(PyJWS):
-    header_type = "JWT"
+class PyJWT:
+    def __init__(self, options=None):
+        if options is None:
+            options = {}
+        self.options = {**self._get_default_options(), **options}
 
     @staticmethod
     def _get_default_options() -> Dict[str, Union[bool, List[str]]]:
@@ -33,7 +36,7 @@ def _get_default_options() -> Dict[str, Union[bool, List[str]]]:
 
     def encode(
         self,
-        payload: Union[Dict, bytes],
+        payload: Dict[str, Any],
         key: str,
         algorithm: str = "HS256",
         headers: Optional[Dict] = None,
@@ -59,20 +62,18 @@ def encode(
             payload, separators=(",", ":"), cls=json_encoder
         ).encode("utf-8")
 
-        return super().encode(
+        return api_jws.encode(
             json_payload, key, algorithm, headers, json_encoder
         )
 
-    def decode(
+    def decode_complete(
         self,
         jwt: str,
         key: str = "",
         algorithms: List[str] = None,
         options: Dict = None,
-        complete: bool = False,
         **kwargs,
     ) -> Dict[str, Any]:
-
         if options is None:
             options = {"verify_signature": True}
         else:
@@ -83,20 +84,16 @@ def decode(
                 'It is required that you pass in a value for the "algorithms" argument when calling decode().'
             )
 
-        decoded = super().decode(
+        decoded = api_jws.decode_complete(
             jwt,
             key=key,
             algorithms=algorithms,
             options=options,
-            complete=complete,
             **kwargs,
         )
 
         try:
-            if complete:
-                payload = json.loads(decoded["payload"])
-            else:
-                payload = json.loads(decoded)
+            payload = json.loads(decoded["payload"])
         except ValueError as e:
             raise DecodeError("Invalid payload string: %s" % e)
         if not isinstance(payload, dict):
@@ -106,11 +103,19 @@ def decode(
             merged_options = {**self.options, **options}
             self._validate_claims(payload, merged_options, **kwargs)
 
-        if complete:
-            decoded["payload"] = payload
-            return decoded
+        decoded["payload"] = payload
+        return decoded
 
-        return payload
+    def decode(
+        self,
+        jwt: str,
+        key: str = "",
+        algorithms: List[str] = None,
+        options: Dict = None,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        decoded = self.decode_complete(jwt, key, algorithms, options, **kwargs)
+        return decoded["payload"]
 
     def _validate_claims(
         self, payload, options, audience=None, issuer=None, leeway=0, **kwargs
@@ -215,7 +220,5 @@ def _validate_iss(self, payload, issuer):
 
 _jwt_global_obj = PyJWT()
 encode = _jwt_global_obj.encode
+decode_complete = _jwt_global_obj.decode_complete
 decode = _jwt_global_obj.decode
-register_algorithm = _jwt_global_obj.register_algorithm
-unregister_algorithm = _jwt_global_obj.unregister_algorithm
-get_unverified_header = _jwt_global_obj.get_unverified_header

jwt/jwks_client.py

@@ -2,7 +2,7 @@
 import urllib.request
 
 from .api_jwk import PyJWKSet
-from .api_jwt import decode as decode_token
+from .api_jwt import decode_complete as decode_token
 from .exceptions import PyJWKClientError
 
 
@@ -50,8 +50,6 @@ def get_signing_key(self, kid):
         return signing_key
 
     def get_signing_key_from_jwt(self, token):
-        unverified = decode_token(
-            token, complete=True, options={"verify_signature": False}
-        )
+        unverified = decode_token(token, options={"verify_signature": False})
         header = unverified["header"]
         return self.get_signing_key(header.get("kid"))

jwt/utils.py

@@ -32,7 +32,7 @@ def base64url_decode(input):
     return base64.urlsafe_b64decode(input)
 
 
-def base64url_encode(input):
+def base64url_encode(input: bytes) -> bytes:
     return base64.urlsafe_b64encode(input).replace(b"=", b"")
 
 

tests/test_api_jws.py

@@ -215,6 +215,27 @@ def test_decodes_valid_jws(self, jws, payload):
 
         assert decoded_payload == payload
 
+    def test_decodes_complete_valid_jws(self, jws, payload):
+        example_secret = "secret"
+        example_jws = (
+            b"eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9."
+            b"aGVsbG8gd29ybGQ."
+            b"gEW0pdU4kxPthjtehYdhxB9mMOGajt1xCKlGGXDJ8PM"
+        )
+
+        decoded = jws.decode_complete(
+            example_jws, example_secret, algorithms=["HS256"]
+        )
+
+        assert decoded == {
+            "header": {"alg": "HS256", "typ": "JWT"},
+            "payload": payload,
+            "signature": (
+                b"\x80E\xb4\xa5\xd58\x93\x13\xed\x86;^\x85\x87a\xc4"
+                b"\x1ff0\xe1\x9a\x8e\xddq\x08\xa9F\x19p\xc9\xf0\xf3"
+            ),
+        }
+
     # 'Control' Elliptic Curve jws created by another library.
     # Used to test for regressions that could affect both
     # encoding / decoding operations equally (causing tests

tests/test_api_jwt.py

@@ -47,6 +47,27 @@ def test_decodes_valid_jwt(self, jwt):
 
         assert decoded_payload == example_payload
 
+    def test_decodes_complete_valid_jwt(self, jwt):
+        example_payload = {"hello": "world"}
+        example_secret = "secret"
+        example_jwt = (
+            b"eyJhbGciOiAiSFMyNTYiLCAidHlwIjogIkpXVCJ9"
+            b".eyJoZWxsbyI6ICJ3b3JsZCJ9"
+            b".tvagLDLoaiJKxOKqpBXSEGy7SYSifZhjntgm9ctpyj8"
+        )
+        decoded = jwt.decode_complete(
+            example_jwt, example_secret, algorithms=["HS256"]
+        )
+
+        assert decoded == {
+            "header": {"alg": "HS256", "typ": "JWT"},
+            "payload": example_payload,
+            "signature": (
+                b'\xb6\xf6\xa0,2\xe8j"J\xc4\xe2\xaa\xa4\x15\xd2'
+                b"\x10l\xbbI\x84\xa2}\x98c\x9e\xd8&\xf5\xcbi\xca?"
+            ),
+        }
+
     def test_load_verify_valid_jwt(self, jwt):
         example_payload = {"hello": "world"}
         example_secret = "secret"
@@ -313,13 +334,12 @@ def test_decode_with_expiration_with_leeway(self, jwt, payload):
         secret = "secret"
         jwt_message = jwt.encode(payload, secret)
 
-        decoded_payload, signing, header, signature = jwt._load(jwt_message)
-
         # With 3 seconds leeway, should be ok
         for leeway in (3, timedelta(seconds=3)):
-            jwt.decode(
+            decoded = jwt.decode(
                 jwt_message, secret, leeway=leeway, algorithms=["HS256"]
             )
+            assert decoded == payload
 
         # With 1 seconds, should fail
         for leeway in (1, timedelta(seconds=1)):