Fix a bug where a PEM private key as bytes raises a TypeError

Author: mark-adams

jwt/algorithms.py

@@ -1,7 +1,7 @@
 import hashlib
 import hmac
 
-from .compat import constant_time_compare, string_types, text_type
+from .compat import binary_type, constant_time_compare, is_string_type
 from .exceptions import InvalidKeyError
 from .utils import der_to_raw_signature, raw_to_der_signature
 
@@ -112,10 +112,10 @@ def __init__(self, hash_alg):
         self.hash_alg = hash_alg
 
     def prepare_key(self, key):
-        if not isinstance(key, string_types) and not isinstance(key, bytes):
+        if not is_string_type(key):
             raise TypeError('Expecting a string- or bytes-formatted key.')
 
-        if isinstance(key, text_type):
+        if not isinstance(key, binary_type):
             key = key.encode('utf-8')
 
         invalid_strings = [
@@ -156,8 +156,8 @@ def prepare_key(self, key):
                isinstance(key, RSAPublicKey):
                 return key
 
-            if isinstance(key, string_types):
-                if isinstance(key, text_type):
+            if is_string_type(key):
+                if not isinstance(key, binary_type):
                     key = key.encode('utf-8')
 
                 try:
@@ -213,8 +213,8 @@ def prepare_key(self, key):
                isinstance(key, EllipticCurvePublicKey):
                 return key
 
-            if isinstance(key, string_types):
-                if isinstance(key, text_type):
+            if is_string_type(key):
+                if not isinstance(key, binary_type):
                     key = key.encode('utf-8')
 
                 # Attempt to load key. We don't know if it's

jwt/compat.py

@@ -11,14 +11,18 @@
 
 
 if PY3:
-    string_types = str,
     text_type = str
     binary_type = bytes
 else:
-    string_types = basestring,
     text_type = unicode
     binary_type = str
 
+string_types = (text_type, binary_type)
+
+
+def is_string_type(val):
+    return any([isinstance(val, typ) for typ in string_types])
+
 
 def timedelta_total_seconds(delta):
     try:

setup.cfg

@@ -1,6 +1,8 @@
 [flake8]
 max-line-length = 119
-exclude = docs/
+exclude = 
+    docs/,
+    .tox/
 
 [wheel]
 universal = 1

tests/test_algorithms.py

@@ -91,6 +91,13 @@ def test_rsa_should_parse_pem_public_key(self):
         with open(key_path('testkey2_rsa.pub.pem'), 'r') as pem_key:
             algo.prepare_key(pem_key.read())
 
+    @pytest.mark.skipif(not has_crypto, reason='Not supported without cryptography library')
+    def test_rsa_should_accept_pem_private_key_bytes(self):
+        algo = RSAAlgorithm(RSAAlgorithm.SHA256)
+
+        with open(key_path('testkey_rsa'), 'rb') as pem_key:
+            algo.prepare_key(pem_key.read())
+
     @pytest.mark.skipif(not has_crypto, reason='Not supported without cryptography library')
     def test_rsa_should_accept_unicode_key(self):
         algo = RSAAlgorithm(RSAAlgorithm.SHA256)
@@ -141,6 +148,13 @@ def test_ec_should_accept_unicode_key(self):
         with open(key_path('testkey_ec'), 'r') as ec_key:
             algo.prepare_key(ensure_unicode(ec_key.read()))
 
+    @pytest.mark.skipif(not has_crypto, reason='Not supported without cryptography library')
+    def test_ec_should_accept_pem_private_key_bytes(self):
+        algo = ECAlgorithm(ECAlgorithm.SHA256)
+
+        with open(key_path('testkey_ec'), 'rb') as ec_key:
+            algo.prepare_key(ec_key.read())
+
     @pytest.mark.skipif(not has_crypto, reason='Not supported without cryptography library')
     def test_ec_verify_should_return_false_if_signature_invalid(self):
         algo = ECAlgorithm(ECAlgorithm.SHA256)