feat: add class method JWSRegistry.guess_alg

ref #49

Author: lepture

src/joserfc/_rfc7515/model.py

@@ -112,6 +112,7 @@ class JWSAlgModel(object, metaclass=ABCMeta):
     key_type = "oct"
     algorithm_type: Literal["JWS"] = "JWS"
     algorithm_location = "sig"
+    algorithm_security = 0
 
     def check_key(self, key: Any) -> None:
         key.check_use("sig")

src/joserfc/_rfc7515/registry.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 import warnings
+from typing import Any
+from enum import Enum
 from .model import JWSAlgModel
-from ..errors import UnsupportedAlgorithmError, SecurityWarning
+from ..errors import UnsupportedAlgorithmError, SecurityWarning, JoseError
 from ..registry import (
     JWS_HEADER_REGISTRY,
     Header,
@@ -28,6 +30,10 @@ class JWSRegistry:
     :param strict_check_header: only allow header key in the registry to be used
     """
 
+    class Strategy(Enum):
+        RECOMMENDED = 1
+        SECURITY = 2
+
     default_header_registry: HeaderRegistryDict = JWS_HEADER_REGISTRY
     algorithms: dict[str, JWSAlgModel] = {}
     recommended: list[str] = []
@@ -79,6 +85,45 @@ def check_header(self, header: Header) -> None:
         if self.strict_check_header:
             check_supported_header(self.header_registry, header)
 
+    @classmethod
+    def guess_alg(cls, key: Any, strategy: Strategy) -> str | None:
+        """Guess the JWS algorithm for a given key.
+
+        :param key: key instance
+        :param strategy: the strategy for guessing the JWS algorithm
+        """
+        if strategy == cls.Strategy.RECOMMENDED:
+            algorithms = cls.filter_algorithms(key, cls.recommended)
+        elif strategy == cls.Strategy.SECURITY:
+            names = list(cls.algorithms.keys())
+            algorithms = cls.filter_algorithms(key, names)
+            # sort by security level
+            algorithms.sort(key=lambda alg: alg.algorithm_security, reverse=True)
+        else:
+            raise NotImplementedError(f"Unknown algorithm strategy '{strategy}'")
+
+        if algorithms:
+            return algorithms[0].name
+        else:
+            return None
+
+    @classmethod
+    def filter_algorithms(cls, key: Any, names: list[str]) -> list[JWSAlgModel]:
+        """Filter JWS algorithms based on the given algorithm names.
+
+        :param key: key instance
+        :param names: list of algorithm names
+        """
+        rv: list[JWSAlgModel] = []
+        for name in names:
+            alg = cls.algorithms[name]
+            try:
+                alg.check_key(key)
+                rv.append(alg)
+            except JoseError:
+                pass
+        return rv
+
 
 #: default JWS registry
 default_registry = JWSRegistry()

src/joserfc/_rfc7518/jws_algs.py

@@ -57,6 +57,7 @@ def __init__(self, sha_type: t.Literal[256, 384, 512], recommended: bool = False
         self.description = f"HMAC using SHA-{sha_type}"
         self.recommended = recommended
         self.hash_alg = getattr(self, f"SHA{sha_type}")
+        self.algorithm_security = sha_type
 
     def sign(self, msg: bytes, key: OctKey) -> bytes:
         # it is faster than the one in cryptography
@@ -89,6 +90,7 @@ def __init__(self, sha_type: t.Literal[256, 384, 512], recommended: bool = False
         self.description = f"RSASSA-PKCS1-v1_5 using SHA-{sha_type}"
         self.recommended = recommended
         self.hash_alg = getattr(self, f"SHA{sha_type}")
+        self.algorithm_security = sha_type
 
     def sign(self, msg: bytes, key: RSAKey) -> bytes:
         op_key = key.get_op_key("sign")
@@ -103,7 +105,7 @@ def verify(self, msg: bytes, sig: bytes, key: RSAKey) -> bool:
             return False
 
 
-class ECAlgorithm(JWSAlgModel):
+class ESAlgorithm(JWSAlgModel):
     """ECDSA using SHA algorithms for JWS. Available algorithms:
 
     - ES256: ECDSA using P-256 and SHA-256
@@ -123,6 +125,7 @@ def __init__(self, name: str, curve: str, sha_type: t.Literal[256, 384, 512], re
         self.description = f"ECDSA using {self.curve} and SHA-{sha_type}"
         self.recommended = recommended
         self.hash_alg = getattr(self, f"SHA{sha_type}")
+        self.algorithm_security = sha_type
 
     def check_key(self, key: ECKey) -> None:
         super().check_key(key)
@@ -174,6 +177,7 @@ def __init__(self, sha_type: t.Literal[256, 384, 512]):
         self.description = f"RSASSA-PSS using SHA-{sha_type} and MGF1 with SHA-{sha_type}"
         self.hash_alg = getattr(self, f"SHA{sha_type}")
         self.padding = padding.PSS(mgf=padding.MGF1(self.hash_alg()), salt_length=self.hash_alg.digest_size)
+        self.algorithm_security = sha_type
 
     def sign(self, msg: bytes, key: RSAKey) -> bytes:
         op_key = key.get_op_key("sign")
@@ -196,9 +200,9 @@ def verify(self, msg: bytes, sig: bytes, key: RSAKey) -> bool:
     RSAAlgorithm(256, True),  # RS256
     RSAAlgorithm(384),  # RS384
     RSAAlgorithm(512),  # RS512
-    ECAlgorithm("ES256", "P-256", 256, True),
-    ECAlgorithm("ES384", "P-384", 384),
-    ECAlgorithm("ES512", "P-521", 512),
+    ESAlgorithm("ES256", "P-256", 256, True),
+    ESAlgorithm("ES384", "P-384", 384),
+    ESAlgorithm("ES512", "P-521", 512),
     RSAPSSAlgorithm(256),  # PS256
     RSAPSSAlgorithm(384),  # PS384
     RSAPSSAlgorithm(512),  # PS512
@@ -208,5 +212,5 @@ def verify(self, msg: bytes, sig: bytes, key: RSAKey) -> bool:
 NoneAlgModel = NoneAlgorithm
 HMACAlgModel = HMACAlgorithm
 RSAAlgModel = RSAAlgorithm
-ECAlgModel = ECAlgorithm
+ECAlgModel = ESAlgorithm
 RSAPSSAlgModel = RSAPSSAlgorithm

src/joserfc/_rfc8812/__init__.py

@@ -1,8 +1,8 @@
 from cryptography.hazmat.primitives.asymmetric.ec import SECP256K1
 from .._rfc7518.ec_key import ECKey
-from .._rfc7518.jws_algs import ECAlgorithm
+from .._rfc7518.jws_algs import ESAlgorithm
 
-ES256K = ECAlgorithm("ES256K", "secp256k1", 256)
+ES256K = ESAlgorithm("ES256K", "secp256k1", 256)
 
 
 def register_secp256k1() -> None:

src/joserfc/jwa.py

@@ -14,7 +14,7 @@
     NoneAlgorithm,
     HMACAlgorithm,
     RSAAlgorithm,
-    ECAlgorithm,
+    ESAlgorithm,
     RSAPSSAlgorithm,
     JWS_ALGORITHMS as _JWS_ALGORITHMS,
 )
@@ -46,7 +46,7 @@
     "NoneAlgorithm",
     "HMACAlgorithm",
     "RSAAlgorithm",
-    "ECAlgorithm",
+    "ESAlgorithm",
     "RSAPSSAlgorithm",
     "EdDSAAlgorithm",
     # JWE algorithms

tests/jws/test_registry.py

@@ -0,0 +1,41 @@
+import unittest
+
+from joserfc.jws import JWSRegistry
+from joserfc.jwk import OctKey, RSAKey, ECKey, OKPKey
+
+
+class JWSRegistryTest(unittest.TestCase):
+    oct_key = OctKey.generate_key()
+    rsa_key = RSAKey.generate_key()
+    ec_key = ECKey.generate_key()
+    okp_key = OKPKey.generate_key()
+
+    def test_guess_recommended_algorithm(self):
+        name = JWSRegistry.guess_alg(self.oct_key, JWSRegistry.Strategy.RECOMMENDED)
+        self.assertEqual(name, "HS256")
+
+        name = JWSRegistry.guess_alg(self.rsa_key, JWSRegistry.Strategy.RECOMMENDED)
+        self.assertEqual(name, "RS256")
+
+        name = JWSRegistry.guess_alg(self.ec_key, JWSRegistry.Strategy.RECOMMENDED)
+        self.assertEqual(name, "ES256")
+
+        name = JWSRegistry.guess_alg(self.okp_key, JWSRegistry.Strategy.RECOMMENDED)
+        self.assertEqual(name, None)
+
+    def test_guess_security_algorithm(self):
+        name = JWSRegistry.guess_alg(self.oct_key, JWSRegistry.Strategy.SECURITY)
+        self.assertEqual(name, "HS512")
+
+        name = JWSRegistry.guess_alg(self.rsa_key, JWSRegistry.Strategy.SECURITY)
+        self.assertEqual(name, "RS512")
+
+        name = JWSRegistry.guess_alg(self.ec_key, JWSRegistry.Strategy.SECURITY)
+        self.assertEqual(name, "ES256")
+
+        ec521 = ECKey.generate_key("P-521")
+        name = JWSRegistry.guess_alg(ec521, JWSRegistry.Strategy.SECURITY)
+        self.assertEqual(name, "ES512")
+
+        name = JWSRegistry.guess_alg(self.okp_key, JWSRegistry.Strategy.SECURITY)
+        self.assertEqual(name, "EdDSA")