Merge pull request #213 from jpadilla/fix-type-error-on-bytes-key

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

3 files changed, 27 insertions, 9 deletions

diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index 9c1a7e8..51e8f16 100644
--- a/jwt/algorithms.py
+++ b/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 @@ class HMACAlgorithm(Algorithm):
         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 @@ if has_crypto:
                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 @@ if has_crypto:
                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
diff --git a/jwt/compat.py b/jwt/compat.py
index 8f6bfed..dafd0c7 100644
--- a/jwt/compat.py
+++ b/jwt/compat.py
@@ -11,14 +11,18 @@ PY3 = sys.version_info[0] == 3
 
 
 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:
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 45bb6d1..e3cf1d0 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -92,6 +92,13 @@ class TestAlgorithms:
             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)
 
@@ -142,6 +149,13 @@ class TestAlgorithms:
             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)