diff --git a/jwt/utils.py b/jwt/utils.py index 81c5ee41..d469139b 100644 --- a/jwt/utils.py +++ b/jwt/utils.py @@ -131,26 +131,15 @@ def is_pem_format(key: bytes) -> bool: # Based on https://github.com/pyca/cryptography/blob/bcb70852d577b3f490f015378c75cba74986297b/src/cryptography/hazmat/primitives/serialization/ssh.py#L40-L46 -_CERT_SUFFIX = b"-cert-v01@openssh.com" -_SSH_PUBKEY_RC = re.compile(rb"\A(\S+)[ \t]+(\S+)") -_SSH_KEY_FORMATS = [ +_SSH_KEY_FORMATS = ( b"ssh-ed25519", b"ssh-rsa", b"ssh-dss", b"ecdsa-sha2-nistp256", b"ecdsa-sha2-nistp384", b"ecdsa-sha2-nistp521", -] +) def is_ssh_key(key: bytes) -> bool: - if any(string_value in key for string_value in _SSH_KEY_FORMATS): - return True - - ssh_pubkey_match = _SSH_PUBKEY_RC.match(key) - if ssh_pubkey_match: - key_type = ssh_pubkey_match.group(1) - if _CERT_SUFFIX == key_type[-len(_CERT_SUFFIX) :]: - return True - - return False + return key.startswith(_SSH_KEY_FORMATS) diff --git a/tests/test_utils.py b/tests/test_utils.py index 122dcb4e..d8d0d6c4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,6 +1,6 @@ import pytest -from jwt.utils import force_bytes, from_base64url_uint, to_base64url_uint +from jwt.utils import force_bytes, from_base64url_uint, is_ssh_key, to_base64url_uint @pytest.mark.parametrize( @@ -37,3 +37,19 @@ def test_from_base64url_uint(inputval, expected): def test_force_bytes_raises_error_on_invalid_object(): with pytest.raises(TypeError): force_bytes({}) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "key_format", + ( + b"ssh-ed25519", + b"ssh-rsa", + b"ssh-dss", + b"ecdsa-sha2-nistp256", + b"ecdsa-sha2-nistp384", + b"ecdsa-sha2-nistp521", + ), +) +def test_is_ssh_key(key_format): + assert is_ssh_key(key_format + b" any") is True + assert is_ssh_key(b"not a ssh key") is False