Add support for PS* algorithms with PyCryptodome

2 files changed, 77 insertions, 3 deletions

diff --git a/jwt/contrib/algorithms/pycryptodome.py b/jwt/contrib/algorithms/pycryptodome.py
index ba5d538..84b1c88 100644
--- a/jwt/contrib/algorithms/pycryptodome.py
+++ b/jwt/contrib/algorithms/pycryptodome.py
@@ -2,7 +2,7 @@ import Cryptodome.Hash.SHA256
 import Cryptodome.Hash.SHA384
 import Cryptodome.Hash.SHA512
 from Cryptodome.PublicKey import ECC, RSA
-from Cryptodome.Signature import DSS, PKCS1_v1_5
+from Cryptodome.Signature import DSS, PKCS1_v1_5, pss
 
 from jwt.algorithms import Algorithm
 from jwt.compat import string_types, text_type
@@ -61,7 +61,6 @@ class ECAlgorithm(Algorithm):
         self.hash_alg = hash_alg
 
     def prepare_key(self, key):
-
         if isinstance(key, ECC.EccKey):
             return key
 
@@ -88,3 +87,39 @@ class ECAlgorithm(Algorithm):
             return True
         except ValueError:
             return False
+
+
+class RSAPSSAlgorithm(RSAAlgorithm):
+    """
+    Performs a signature using RSASSA-PSS with MGF1
+
+    This class requires the PyCryptodome package to be installed.
+    """
+
+    def prepare_key(self, key):
+        if isinstance(key, ECC.EccKey):
+            return key
+
+        if isinstance(key, string_types):
+            if isinstance(key, text_type):
+                key = key.encode("utf-8")
+            key = RSA.import_key(key)
+        else:
+            raise TypeError("Expecting a PEM- or RSA-formatted key.")
+
+        return key
+
+    def sign(self, msg, key):
+        signer = pss.new(key)
+        hash_obj = self.hash_alg.new(msg)
+        return signer.sign(hash_obj)
+
+    def verify(self, msg, key, sig):
+        hash_obj = self.hash_alg.new(msg)
+        verifier = pss.new(key)
+
+        try:
+            verifier.verify(hash_obj, sig)
+            return True
+        except (ValueError, TypeError):
+            return False
diff --git a/tests/contrib/test_algorithms.py b/tests/contrib/test_algorithms.py
index 17d8f94..b565a1b 100644
--- a/tests/contrib/test_algorithms.py
+++ b/tests/contrib/test_algorithms.py
@@ -22,7 +22,7 @@ except ImportError:
 
 try:
     # fmt: off
-    from jwt.contrib.algorithms.pycryptodome import RSAAlgorithm, ECAlgorithm  # noqa: F811
+    from jwt.contrib.algorithms.pycryptodome import RSAAlgorithm, ECAlgorithm, RSAPSSAlgorithm  # noqa: F811
     # fmt: on
 
     has_pycryptodome = True
@@ -406,3 +406,42 @@ class TestPyCryptodomeAlgorithms:
             jwt_pub_key_second = algo.prepare_key(jwt_pub_key_first)
 
         assert jwt_pub_key_first == jwt_pub_key_second
+
+    def test_rsa_pss_sign_then_verify_should_return_true(self):
+        algo = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256)
+
+        message = force_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)
+        assert result
+
+    def test_rsa_pss_verify_should_return_false_if_signature_invalid(self):
+        algo = RSAPSSAlgorithm(RSAPSSAlgorithm.SHA256)
+
+        jwt_message = force_bytes("Hello World!")
+
+        jwt_sig = base64.b64decode(
+            force_bytes(
+                "ywKAUGRIDC//6X+tjvZA96yEtMqpOrSppCNfYI7NKyon3P7doud5v65oWNu"
+                "vQsz0fzPGfF7mQFGo9Cm9Vn0nljm4G6PtqZRbz5fXNQBH9k10gq34AtM02c"
+                "/cveqACQ8gF3zxWh6qr9jVqIpeMEaEBIkvqG954E0HT9s9ybHShgHX9mlWk"
+                "186/LopP4xe5c/hxOQjwhv6yDlTiwJFiqjNCvj0GyBKsc4iECLGIIO+4mC4"
+                "daOCWqbpZDuLb1imKpmm8Nsm56kAxijMLZnpCcnPgyb7CqG+B93W9GHglA5"
+                "drUeR1gRtO7vqbZMsCAQ4bpjXxwbYyjQlEVuMl73UL6sOWg=="
+            )
+        )
+
+        jwt_sig += force_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)
+        assert not result