Merge pull request #132 from mark-adams/add-rsa-pss-algorithms

Added support for RSASSA-PSS algorithms (PS256, PS384, PS512)

5 files changed, 124 insertions, 16 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1564f5e..7e2b77d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,8 +7,10 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 [Unreleased][unreleased]
 -------------------------------------------------------------------------
 ### Changed
+- Added flexible and complete verification options during decode #131
+- Added support for PS256, PS384, and PS512 algorithms. #132
 - Added this CHANGELOG.md file
-- Added flexible and complete verification options. #131
+
 
 ### Fixed
 - Placeholder
diff --git a/README.md b/README.md
index 5ae0b40..22d088b 100644
--- a/README.md
+++ b/README.md
@@ -63,7 +63,7 @@ except jwt.InvalidTokenError:
 ```
 
 You may also override exception checking via an `options` dictionary.  The default
-options are as follows: 
+options are as follows:
 
 ```python
 options = {
@@ -112,6 +112,9 @@ currently supports:
 * RS256 - RSASSA-PKCS1-v1_5 signature algorithm using SHA-256 hash algorithm
 * RS384 - RSASSA-PKCS1-v1_5 signature algorithm using SHA-384 hash algorithm
 * RS512 - RSASSA-PKCS1-v1_5 signature algorithm using SHA-512 hash algorithm
+* PS256 - RSASSA-PSS signature using SHA-256 and MGF1 padding with SHA-256  
+* PS384 - RSASSA-PSS signature using SHA-384 and MGF1 padding with SHA-384
+* PS512 - RSASSA-PSS signature using SHA-512 and MGF1 padding with SHA-512
 
 ### Encoding
 You can specify which algorithm you would like to use to sign the JWT
diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index cda7b19..05fd194 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -42,7 +42,10 @@ def get_default_algorithms():
             'RS512': RSAAlgorithm(RSAAlgorithm.SHA512),
             'ES256': ECAlgorithm(ECAlgorithm.SHA256),
             'ES384': ECAlgorithm(ECAlgorithm.SHA384),
-            'ES512': ECAlgorithm(ECAlgorithm.SHA512)
+            'ES512': ECAlgorithm(ECAlgorithm.SHA512),
+            'PS256': RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256),
+            'PS384': RSAPSSAlgorithm(RSAPSSAlgorithm.SHA384),
+            'PS512': RSAPSSAlgorithm(RSAPSSAlgorithm.SHA512)
         })
 
     return default_algorithms
@@ -145,7 +148,7 @@ if has_crypto:
         SHA512 = hashes.SHA512
 
         def __init__(self, hash_alg):
-            self.hash_alg = hash_alg()
+            self.hash_alg = hash_alg
 
         def prepare_key(self, key):
             if isinstance(key, RSAPrivateKey) or \
@@ -171,7 +174,7 @@ if has_crypto:
         def sign(self, msg, key):
             signer = key.signer(
                 padding.PKCS1v15(),
-                self.hash_alg
+                self.hash_alg()
             )
 
             signer.update(msg)
@@ -181,7 +184,7 @@ if has_crypto:
             verifier = key.verifier(
                 sig,
                 padding.PKCS1v15(),
-                self.hash_alg
+                self.hash_alg()
             )
 
             verifier.update(msg)
@@ -202,7 +205,7 @@ if has_crypto:
         SHA512 = hashes.SHA512
 
         def __init__(self, hash_alg):
-            self.hash_alg = hash_alg()
+            self.hash_alg = hash_alg
 
         def prepare_key(self, key):
             if isinstance(key, EllipticCurvePrivateKey) or \
@@ -227,13 +230,48 @@ if has_crypto:
             return key
 
         def sign(self, msg, key):
-            signer = key.signer(ec.ECDSA(self.hash_alg))
+            signer = key.signer(ec.ECDSA(self.hash_alg()))
 
             signer.update(msg)
             return signer.finalize()
 
         def verify(self, msg, key, sig):
-            verifier = key.verifier(sig, ec.ECDSA(self.hash_alg))
+            verifier = key.verifier(sig, ec.ECDSA(self.hash_alg()))
+
+            verifier.update(msg)
+
+            try:
+                verifier.verify()
+                return True
+            except InvalidSignature:
+                return False
+
+    class RSAPSSAlgorithm(RSAAlgorithm):
+        """
+        Performs a signature using RSASSA-PSS with MGF1
+        """
+
+        def sign(self, msg, key):
+            signer = key.signer(
+                padding.PSS(
+                    mgf=padding.MGF1(self.hash_alg()),
+                    salt_length=padding.PSS.MAX_LENGTH
+                ),
+                self.hash_alg()
+            )
+
+            signer.update(msg)
+            return signer.finalize()
+
+        def verify(self, msg, key, sig):
+            verifier = key.verifier(
+                sig,
+                padding.PSS(
+                    mgf=padding.MGF1(self.hash_alg()),
+                    salt_length=padding.PSS.MAX_LENGTH
+                ),
+                self.hash_alg()
+            )
 
             verifier.update(msg)
 
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 3ea9a04..8eb0eb3 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -7,7 +7,7 @@ from .compat import unittest
 from .utils import ensure_bytes, ensure_unicode, key_path
 
 try:
-    from jwt.algorithms import RSAAlgorithm, ECAlgorithm
+    from jwt.algorithms import RSAAlgorithm, ECAlgorithm, RSAPSSAlgorithm
 
     has_crypto = True
 except ImportError:
@@ -169,34 +169,92 @@ class TestAlgorithms(unittest.TestCase):
     def test_ec_verify_should_return_false_if_signature_invalid(self):
         algo = ECAlgorithm(ECAlgorithm.SHA256)
 
-        jwt_message = ensure_bytes('Hello World!')
+        message = ensure_bytes('Hello World!')
 
         # Mess up the signature by replacing a known byte
-        jwt_sig = base64.b64decode(ensure_bytes(
+        sig = base64.b64decode(ensure_bytes(
             'MIGIAkIB9vYz+inBL8aOTA4auYz/zVuig7TT1bQgKROIQX9YpViHkFa4DT5'
             '5FuFKn9XzVlk90p6ldEj42DC9YecXHbC2t+cCQgCicY+8f3f/KCNtWK7cif'
             '6vdsVwm6Lrjs0Ag6ZqCf+olN11hVt1qKBC4lXppqB1gNWEmNQaiz1z2QRyc'
             'zJ8hSJmbw=='.replace('r', 's')))
 
         with open(key_path('testkey_ec.pub'), 'r') as keyfile:
-            jwt_pub_key = algo.prepare_key(keyfile.read())
+            pub_key = algo.prepare_key(keyfile.read())
 
-        result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
+        result = algo.verify(message, pub_key, sig)
         self.assertFalse(result)
 
     @unittest.skipIf(not has_crypto, 'Not supported without cryptography library')
     def test_ec_verify_should_return_true_if_signature_valid(self):
         algo = ECAlgorithm(ECAlgorithm.SHA256)
 
-        jwt_message = ensure_bytes('Hello World!')
+        message = ensure_bytes('Hello World!')
 
-        jwt_sig = base64.b64decode(ensure_bytes(
+        sig = base64.b64decode(ensure_bytes(
             'MIGIAkIB9vYz+inBL8aOTA4auYz/zVuig7TT1bQgKROIQX9YpViHkFa4DT5'
             '5FuFKn9XzVlk90p6ldEj42DC9YecXHbC2t+cCQgCicY+8f3f/KCNtWK7cif'
             '6vdsVwm6Lrjs0Ag6ZqCf+olN11hVt1qKBC4lXppqB1gNWEmNQaiz1z2QRyc'
             'zJ8hSJmbw=='))
 
         with open(key_path('testkey_ec.pub'), 'r') as keyfile:
+            pub_key = algo.prepare_key(keyfile.read())
+
+        result = algo.verify(message, pub_key, sig)
+        self.assertTrue(result)
+
+    @unittest.skipIf(not has_crypto, 'Not supported without cryptography library')
+    def test_rsa_pss_sign_then_verify_should_return_true(self):
+        algo = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256)
+
+        message = ensure_bytes('Hello World!')
+
+        with open(key_path('testkey_rsa'), 'r') as keyfile:
+            priv_key = algo.prepare_key(keyfile.read())
+            sig = algo.sign(message, priv_key)
+
+        with open(key_path('testkey_rsa.pub'), 'r') as keyfile:
+            pub_key = algo.prepare_key(keyfile.read())
+
+        result = algo.verify(message, pub_key, sig)
+        self.assertTrue(result)
+
+    @unittest.skipIf(not has_crypto, 'Not supported without cryptography library')
+    def test_rsa_pss_verify_should_return_false_if_signature_invalid(self):
+        algo = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256)
+
+        jwt_message = ensure_bytes('Hello World!')
+
+        jwt_sig = base64.b64decode(ensure_bytes(
+            'ywKAUGRIDC//6X+tjvZA96yEtMqpOrSppCNfYI7NKyon3P7doud5v65oWNu'
+            'vQsz0fzPGfF7mQFGo9Cm9Vn0nljm4G6PtqZRbz5fXNQBH9k10gq34AtM02c'
+            '/cveqACQ8gF3zxWh6qr9jVqIpeMEaEBIkvqG954E0HT9s9ybHShgHX9mlWk'
+            '186/LopP4xe5c/hxOQjwhv6yDlTiwJFiqjNCvj0GyBKsc4iECLGIIO+4mC4'
+            'daOCWqbpZDuLb1imKpmm8Nsm56kAxijMLZnpCcnPgyb7CqG+B93W9GHglA5'
+            'drUeR1gRtO7vqbZMsCAQ4bpjXxwbYyjQlEVuMl73UL6sOWg=='))
+
+        jwt_sig += ensure_bytes('123')  # Signature is now invalid
+
+        with open(key_path('testkey_rsa.pub'), 'r') as keyfile:
+            jwt_pub_key = algo.prepare_key(keyfile.read())
+
+        result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
+        self.assertFalse(result)
+
+    @unittest.skipIf(not has_crypto, 'Not supported without cryptography library')
+    def test_rsa_pss_verify_should_return_true_if_signature_valid(self):
+        algo = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256)
+
+        jwt_message = ensure_bytes('Hello World!')
+
+        jwt_sig = base64.b64decode(ensure_bytes(
+            'ywKAUGRIDC//6X+tjvZA96yEtMqpOrSppCNfYI7NKyon3P7doud5v65oWNu'
+            'vQsz0fzPGfF7mQFGo9Cm9Vn0nljm4G6PtqZRbz5fXNQBH9k10gq34AtM02c'
+            '/cveqACQ8gF3zxWh6qr9jVqIpeMEaEBIkvqG954E0HT9s9ybHShgHX9mlWk'
+            '186/LopP4xe5c/hxOQjwhv6yDlTiwJFiqjNCvj0GyBKsc4iECLGIIO+4mC4'
+            'daOCWqbpZDuLb1imKpmm8Nsm56kAxijMLZnpCcnPgyb7CqG+B93W9GHglA5'
+            'drUeR1gRtO7vqbZMsCAQ4bpjXxwbYyjQlEVuMl73UL6sOWg=='))
+
+        with open(key_path('testkey_rsa.pub'), 'r') as keyfile:
             jwt_pub_key = algo.prepare_key(keyfile.read())
 
         result = algo.verify(jwt_message, jwt_pub_key, jwt_sig)
diff --git a/tests/test_api.py b/tests/test_api.py
index 33ccd51..2364cc2 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -615,10 +615,17 @@ class TestAPI(unittest.TestCase):
             self.assertTrue('RS256' in jwt_algorithms)
             self.assertTrue('RS384' in jwt_algorithms)
             self.assertTrue('RS512' in jwt_algorithms)
+            self.assertTrue('PS256' in jwt_algorithms)
+            self.assertTrue('PS384' in jwt_algorithms)
+            self.assertTrue('PS512' in jwt_algorithms)
+
         else:
             self.assertFalse('RS256' in jwt_algorithms)
             self.assertFalse('RS384' in jwt_algorithms)
             self.assertFalse('RS512' in jwt_algorithms)
+            self.assertFalse('PS256' in jwt_algorithms)
+            self.assertFalse('PS384' in jwt_algorithms)
+            self.assertFalse('PS512' in jwt_algorithms)
 
     @unittest.skipIf(not has_crypto, "Can't run without cryptography library")
     def test_encode_decode_with_ecdsa_sha256(self):