Merge pull request #23 from jpadilla/patch-1

Make PyCrypto optional if you're not using any of the RSASSA-PKCS1-v1_5 algorithms

3 files changed, 86 insertions, 38 deletions

diff --git a/README.md b/README.md
index 2b84912..33d9a4a 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,10 @@ Installing
 ----------
 
     sudo easy_install PyJWT
+    
+**Note**: The RSASSA-PKCS1-v1_5 algorithms depend on PyCrypto. If you plan on using any of those algorithms you'll need to install it as well.
+
+	sudo easy_install PyCrypto
 
 Usage
 -----
@@ -38,8 +42,7 @@ Change the algorithm with by setting it in encode:
     jwt.encode({"some": "payload"}, "secret", "HS512")
 
 For the RSASSA-PKCS1-v1_5 algorithms, the "secret" argument in jwt.encode is supposed to be a private RSA key as
-imported with Crypto.PublicKey.RSA.importKey. Likewise, the "secret" argument in jwt.decode is supposed to be the
-public RSA key imported with the same method. 
+imported with Crypto.PublicKey.RSA.importKey. Likewise, the "secret" argument in jwt.decode is supposed to be the public RSA key imported with the same method.
 
 Tests
 -----
diff --git a/jwt/__init__.py b/jwt/__init__.py
index b73aa08..708d7ca 100644
--- a/jwt/__init__.py
+++ b/jwt/__init__.py
@@ -14,11 +14,6 @@ from datetime import datetime
 from calendar import timegm
 from collections import Mapping
 
-from Crypto.Signature import PKCS1_v1_5
-from Crypto.Hash import SHA256
-from Crypto.Hash import SHA384
-from Crypto.Hash import SHA512
-
 try:
     import json
 except ImportError:
@@ -42,20 +37,34 @@ class ExpiredSignature(Exception):
 signing_methods = {
     'HS256': lambda msg, key: hmac.new(key, msg, hashlib.sha256).digest(),
     'HS384': lambda msg, key: hmac.new(key, msg, hashlib.sha384).digest(),
-    'HS512': lambda msg, key: hmac.new(key, msg, hashlib.sha512).digest(),
-    'RS256': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA256.new(msg)),
-    'RS384': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA384.new(msg)),
-    'RS512': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA512.new(msg)),
-    }
+    'HS512': lambda msg, key: hmac.new(key, msg, hashlib.sha512).digest()
+}
 
 verify_methods = {
     'HS256': lambda msg, key: hmac.new(key, msg, hashlib.sha256).digest(),
     'HS384': lambda msg, key: hmac.new(key, msg, hashlib.sha384).digest(),
-    'HS512': lambda msg, key: hmac.new(key, msg, hashlib.sha512).digest(),
-    'RS256': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA256.new(msg), sig),
-    'RS384': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA384.new(msg), sig),
-    'RS512': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA512.new(msg), sig),
-    }
+    'HS512': lambda msg, key: hmac.new(key, msg, hashlib.sha512).digest()
+}
+
+try:
+    from Crypto.Signature import PKCS1_v1_5
+    from Crypto.Hash import SHA256
+    from Crypto.Hash import SHA384
+    from Crypto.Hash import SHA512
+    
+    signing_methods.update({
+        'RS256': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA256.new(msg)),
+        'RS384': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA384.new(msg)),
+        'RS512': lambda msg, key: PKCS1_v1_5.new(key).sign(SHA512.new(msg))
+    })
+    
+    verify_methods.update({
+        'RS256': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA256.new(msg), sig),
+        'RS384': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA384.new(msg), sig),
+        'RS512': lambda msg, key, sig: PKCS1_v1_5.new(key).verify(SHA512.new(msg), sig)
+    })
+except ImportError:
+    pass
 
 
 def constant_time_compare(val1, val2):
diff --git a/tests/test_jwt.py b/tests/test_jwt.py
index 6404088..fb83136 100644
--- a/tests/test_jwt.py
+++ b/tests/test_jwt.py
@@ -10,8 +10,6 @@ import jwt
 if sys.version_info >= (3, 0, 0):
     unicode = str
 
-from Crypto.PublicKey import RSA
-
 
 def utc_timestamp():
     return timegm(datetime.utcnow().utctimetuple())
@@ -175,30 +173,68 @@ class TestJWT(unittest.TestCase):
             jwt.decode(jwt_message, secret, leeway=1)
 
     def test_encode_decode_with_rsa_sha256(self):
-        with open('tests/testkey','r') as rsa_priv_file:
-            priv_rsakey = RSA.importKey(rsa_priv_file.read())
-            jwt_message = jwt.encode(self.payload, priv_rsakey, algorithm='RS256')
-        with open('tests/testkey.pub','r') as rsa_pub_file:
-            pub_rsakey = RSA.importKey(rsa_pub_file.read())
-            assert jwt.decode(jwt_message, pub_rsakey)
+        try:
+            from Crypto.PublicKey import RSA
+
+            with open('tests/testkey','r') as rsa_priv_file:
+                priv_rsakey = RSA.importKey(rsa_priv_file.read())
+                jwt_message = jwt.encode(self.payload, priv_rsakey, algorithm='RS256')
+
+            with open('tests/testkey.pub','r') as rsa_pub_file:
+                pub_rsakey = RSA.importKey(rsa_pub_file.read())
+                assert jwt.decode(jwt_message, pub_rsakey)
+        except ImportError:
+            pass
 
     def test_encode_decode_with_rsa_sha384(self):
-        with open('tests/testkey','r') as rsa_priv_file:
-            priv_rsakey = RSA.importKey(rsa_priv_file.read())
-            jwt_message = jwt.encode(self.payload, priv_rsakey, algorithm='RS384')
-        with open('tests/testkey.pub','r') as rsa_pub_file:
-            pub_rsakey = RSA.importKey(rsa_pub_file.read())
-            assert jwt.decode(jwt_message, pub_rsakey)
+        try:
+            from Crypto.PublicKey import RSA
 
-    def test_encode_decode_with_rsa_sha512(self):
-        with open('tests/testkey','r') as rsa_priv_file:
-            priv_rsakey = RSA.importKey(rsa_priv_file.read())
-            jwt_message = jwt.encode(self.payload, priv_rsakey, algorithm='RS512')
-        with open('tests/testkey.pub','r') as rsa_pub_file:
-            pub_rsakey = RSA.importKey(rsa_pub_file.read())
-            assert jwt.decode(jwt_message, pub_rsakey)
+            with open('tests/testkey','r') as rsa_priv_file:
+                priv_rsakey = RSA.importKey(rsa_priv_file.read())
+                jwt_message = jwt.encode(self.payload, priv_rsakey, algorithm='RS384')
 
+            with open('tests/testkey.pub','r') as rsa_pub_file:
+                pub_rsakey = RSA.importKey(rsa_pub_file.read())
+                assert jwt.decode(jwt_message, pub_rsakey)
+        except ImportError:
+            pass
 
+    def test_encode_decode_with_rsa_sha512(self):
+        try:
+            from Crypto.PublicKey import RSA
+
+            with open('tests/testkey','r') as rsa_priv_file:
+                priv_rsakey = RSA.importKey(rsa_priv_file.read())
+                jwt_message = jwt.encode(self.payload, priv_rsakey, algorithm='RS512')
+
+            with open('tests/testkey.pub','r') as rsa_pub_file:
+                pub_rsakey = RSA.importKey(rsa_pub_file.read())
+                assert jwt.decode(jwt_message, pub_rsakey)
+        except ImportError:
+            pass
+
+    def test_crypto_related_signing_methods(self):
+        try:
+            import Crypto
+            self.assertTrue('RS256' in jwt.signing_methods)
+            self.assertTrue('RS384' in jwt.signing_methods)
+            self.assertTrue('RS512' in jwt.signing_methods)
+        except ImportError:
+            self.assertFalse('RS256' in jwt.signing_methods)
+            self.assertFalse('RS384' in jwt.signing_methods)
+            self.assertFalse('RS512' in jwt.signing_methods)
+
+    def test_crypto_related_verify_methods(self):
+        try:
+            import Crypto
+            self.assertTrue('RS256' in jwt.verify_methods)
+            self.assertTrue('RS384' in jwt.verify_methods)
+            self.assertTrue('RS512' in jwt.verify_methods)
+        except ImportError:
+            self.assertFalse('RS256' in jwt.verify_methods)
+            self.assertFalse('RS384' in jwt.verify_methods)
+            self.assertFalse('RS512' in jwt.verify_methods)
 
 
 if __name__ == '__main__':