From adc90d270a0789b6ddedd6a0efde36990b0e0f59 Mon Sep 17 00:00:00 2001
From: Stephen Rosen <sirosen@globus.org>
Date: Sun, 3 Jul 2022 15:44:00 +0000
Subject: [PATCH] Add compute_hash_digest to Algorithm objects

`Algorithm.compute_hash_digest` is defined as a method which inspects
the object to see that it has the requisite attributes, `hash_alg`.

If `hash_alg` is not set, then the method raises a
NotImplementedError. This applies to classes like NoneAlgorithm.

If `hash_alg` is set, then it is checked for
```
has_crypto  # is cryptography available?
and isinstance(hash_alg, type)
and issubclass(hash_alg, hashes.HashAlgorithm)
```
to see which API for computing a digest is appropriate --
`hashlib` vs `cryptography.hazmat.primitives.hashes`.

These checks could be avoided at runtime if it were necessary to
optimize further (e.g. attach compute_hash_digest methods to classes
with a class decorator) but this is not clearly a worthwhile
optimization. Such perf tuning is intentionally omitted for now.
---
 CHANGELOG.rst            |  4 ++++
 jwt/algorithms.py        | 23 +++++++++++++++++++++++
 tests/test_algorithms.py | 23 +++++++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8ec7ede8..4d562070 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -18,6 +18,10 @@ Fixed
 Added
 ~~~~~
 
+- Add ``compute_hash_digest`` as a method of ``Algorithm`` objects, which uses
+  the underlying hash algorithm to compute a digest. If there is no appropriate
+  hash algorithm, a ``NotImplementedError`` will be raised
+
 `v2.6.0 <https://github.com/jpadilla/pyjwt/compare/2.5.0...2.6.0>`__
 -----------------------------------------------------------------------
 
diff --git a/jwt/algorithms.py b/jwt/algorithms.py
index 93fadf4c..4fae441f 100644
--- a/jwt/algorithms.py
+++ b/jwt/algorithms.py
@@ -18,6 +18,7 @@
 try:
     import cryptography.exceptions
     from cryptography.exceptions import InvalidSignature
+    from cryptography.hazmat.backends import default_backend
     from cryptography.hazmat.primitives import hashes
     from cryptography.hazmat.primitives.asymmetric import ec, padding
     from cryptography.hazmat.primitives.asymmetric.ec import (
@@ -111,6 +112,28 @@ class Algorithm:
     The interface for an algorithm used to sign and verify tokens.
     """
 
+    def compute_hash_digest(self, bytestr: bytes) -> bytes:
+        """
+        Compute a hash digest using the specified algorithm's hash algorithm.
+
+        If there is no hash algorithm, raises a NotImplementedError.
+        """
+        # lookup self.hash_alg if defined in a way that mypy can understand
+        hash_alg = getattr(self, "hash_alg", None)
+        if hash_alg is None:
+            raise NotImplementedError
+
+        if (
+            has_crypto
+            and isinstance(hash_alg, type)
+            and issubclass(hash_alg, hashes.HashAlgorithm)
+        ):
+            digest = hashes.Hash(hash_alg(), backend=default_backend())
+            digest.update(bytestr)
+            return digest.finalize()
+        else:
+            return hash_alg(bytestr).digest()
+
     def prepare_key(self, key):
         """
         Performs necessary validation and conversions on the key and returns
diff --git a/tests/test_algorithms.py b/tests/test_algorithms.py
index 538078af..894ce282 100644
--- a/tests/test_algorithms.py
+++ b/tests/test_algorithms.py
@@ -45,6 +45,12 @@ def test_algorithm_should_throw_exception_if_from_jwk_not_impl(self):
         with pytest.raises(NotImplementedError):
             algo.to_jwk("value")
 
+    def test_algorithm_should_throw_exception_if_compute_hash_digest_not_impl(self):
+        algo = Algorithm()
+
+        with pytest.raises(NotImplementedError):
+            algo.compute_hash_digest(b"value")
+
     def test_none_algorithm_should_throw_exception_if_key_is_not_none(self):
         algo = NoneAlgorithm()
 
@@ -1054,3 +1060,20 @@ def test_okp_ed448_to_jwk_works_with_from_jwk(self):
         signature_2 = algo.sign(b"Hello World!", priv_key_2)
         assert algo.verify(b"Hello World!", pub_key_2, signature_1)
         assert algo.verify(b"Hello World!", pub_key_2, signature_2)
+
+    @crypto_required
+    def test_rsa_can_compute_digest(self):
+        # this is the well-known sha256 hash of "foo"
+        foo_hash = base64.b64decode(b"LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=")
+
+        algo = RSAAlgorithm(RSAAlgorithm.SHA256)
+        computed_hash = algo.compute_hash_digest(b"foo")
+        assert computed_hash == foo_hash
+
+    def test_hmac_can_compute_digest(self):
+        # this is the well-known sha256 hash of "foo"
+        foo_hash = base64.b64decode(b"LCa0a2j/xo/5m0U8HTBBNBNCLXBkg7+g+YpeiGJm564=")
+
+        algo = HMACAlgorithm(HMACAlgorithm.SHA256)
+        computed_hash = algo.compute_hash_digest(b"foo")
+        assert computed_hash == foo_hash