Skip to content

Commit

Permalink
Handle load_pem_public_key ValueError (#952)
Browse files Browse the repository at this point in the history
* Handle load_pem_public_key ValueError

* Add test for invalid key errors on prepare_key of an invalid key

---------

Co-authored-by: MVRA <61065@icf.com>
  • Loading branch information
CollinEMac and MVRA authored Jun 10, 2024
1 parent 9dc732f commit 4703f87
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 2 deletions.
7 changes: 5 additions & 2 deletions jwt/algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
)

try:
from cryptography.exceptions import InvalidSignature
from cryptography.exceptions import InvalidSignature, UnsupportedAlgorithm
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
Expand Down Expand Up @@ -343,7 +343,10 @@ def prepare_key(self, key: AllowedRSAKeys | str | bytes) -> AllowedRSAKeys:
RSAPrivateKey, load_pem_private_key(key_bytes, password=None)
)
except ValueError:
return cast(RSAPublicKey, load_pem_public_key(key_bytes))
try:
return cast(RSAPublicKey, load_pem_public_key(key_bytes))
except (ValueError, UnsupportedAlgorithm):
raise InvalidKeyError("Could not parse the provided public key.")

@overload
@staticmethod
Expand Down
11 changes: 11 additions & 0 deletions tests/test_algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -1100,3 +1100,14 @@ def test_hmac_can_compute_digest(self):
algo = HMACAlgorithm(HMACAlgorithm.SHA256)
computed_hash = algo.compute_hash_digest(b"foo")
assert computed_hash == foo_hash

@crypto_required
def test_rsa_prepare_key_raises_invalid_key_error_on_invalid_pem(self):
algo = RSAAlgorithm(RSAAlgorithm.SHA256)
invalid_key = "invalid key"

with pytest.raises(InvalidKeyError) as excinfo:
algo.prepare_key(invalid_key)

# Check that the exception message is correct
assert "Could not parse the provided public key." in str(excinfo.value)

0 comments on commit 4703f87

Please sign in to comment.