From a0f02777f0bfd71b90e54c0ca42ea3745e925066 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 20 Jan 2024 11:51:38 -1000 Subject: [PATCH 1/2] Add coverage and improve performance of is_ssh_key --- jwt/utils.py | 18 ++++++++++-------- tests/test_utils.py | 8 +++++++- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/jwt/utils.py b/jwt/utils.py index 81c5ee41..f83ceba5 100644 --- a/jwt/utils.py +++ b/jwt/utils.py @@ -133,24 +133,26 @@ 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): + if key.startswith(_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 + # Avoid the regex if the _CERT_SUFFIX is not in the key + if _CERT_SUFFIX in key: + 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 diff --git a/tests/test_utils.py b/tests/test_utils.py index 122dcb4e..a83aff02 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,9 @@ 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] + + +def test_is_ssh_key(): + assert is_ssh_key(b"ecdsa-sha2-nistp256 any") is True + assert is_ssh_key(b"not a ssh key") is False + assert is_ssh_key(b"any-cert-v01@openssh.com any") is True From 29ea595541d0ff9af05b8683347167ea4dd8ae6d Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 20 Jan 2024 12:10:51 -1000 Subject: [PATCH 2/2] simplify --- jwt/utils.py | 15 +-------------- tests/test_utils.py | 16 +++++++++++++--- 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/jwt/utils.py b/jwt/utils.py index f83ceba5..d469139b 100644 --- a/jwt/utils.py +++ b/jwt/utils.py @@ -131,8 +131,6 @@ 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 = ( b"ssh-ed25519", b"ssh-rsa", @@ -144,15 +142,4 @@ def is_pem_format(key: bytes) -> bool: def is_ssh_key(key: bytes) -> bool: - if key.startswith(_SSH_KEY_FORMATS): - return True - - # Avoid the regex if the _CERT_SUFFIX is not in the key - if _CERT_SUFFIX in key: - 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 a83aff02..d8d0d6c4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -39,7 +39,17 @@ def test_force_bytes_raises_error_on_invalid_object(): force_bytes({}) # type: ignore[arg-type] -def test_is_ssh_key(): - assert is_ssh_key(b"ecdsa-sha2-nistp256 any") is True +@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 - assert is_ssh_key(b"any-cert-v01@openssh.com any") is True