feat: filter_algorithms supports KeySet objects

Author: azmeuk

docs/changelog.rst

@@ -19,6 +19,7 @@ Changelog
 
 - ``filter_algorithms`` ``names`` defaults to all algorithms. :pr:`79`.
 - Replace ``JWSRegistry.guess_alg`` with ``JWSRegistry.guess_algorithm``.
+- ``filter_algorithms`` and ``guess_alg`` supports ``KeySet`` objects. :pr:`81`
 
 1.5.0
 -----

src/joserfc/_rfc7515/registry.py

@@ -17,6 +17,7 @@
     check_crit_header,
     check_supported_header,
 )
+from .._keys import KeySet
 
 __all__ = [
     "JWSRegistry",
@@ -115,7 +116,7 @@ def validate_signature_size(self, signature: bytes) -> None:
     def guess_algorithm(cls, key: Any, strategy: Strategy) -> JWSAlgModel | None:
         """Guess the JWS algorithm for a given key.
 
-        :param key: key instance
+        :param key: key instance or a KeySet
         :param strategy: the strategy for guessing the JWS algorithm
         """
         if strategy == cls.Strategy.RECOMMENDED:
@@ -145,12 +146,19 @@ def guess_alg(cls, key: Any, strategy: Strategy) -> str | None:  # pragma: no co
     def filter_algorithms(cls, key: Any, names: list[str] | None = None) -> list[JWSAlgModel]:
         """Filter JWS algorithms based on the given algorithm names.
 
-        :param key: key instance
+        :param key: a key instance or a KeySet
         :param names: list of algorithm names
         """
         if names is None:
             names = list(cls.algorithms.keys())
         rv: list[JWSAlgModel] = []
+        if isinstance(key, KeySet):
+            for k in key.keys:
+                for alg in cls.filter_algorithms(k, names):
+                    if alg not in rv:
+                        rv.append(alg)
+            return rv
+
         for name in names:
             alg = cls.algorithms[name]
             try:

tests/jws/test_registry.py

@@ -1,7 +1,7 @@
 import unittest
 
 from joserfc.jws import JWSRegistry
-from joserfc.jwk import OctKey, RSAKey, ECKey, OKPKey
+from joserfc.jwk import OctKey, RSAKey, ECKey, OKPKey, KeySet
 
 
 class JWSRegistryTest(unittest.TestCase):
@@ -63,3 +63,32 @@ def test_filter_algorithms_ed448(self):
         self.assertIn("EdDSA", names)
         self.assertIn("Ed448", names)
         self.assertNotIn("Ed25519", names)
+
+    def test_filter_algorithms_with_key_set(self):
+        """filter_algorithms should support KeySet and combine algorithms from all keys."""
+        rsa_key1 = RSAKey.generate_key()
+        rsa_key2 = RSAKey.generate_key()
+        ec_key = ECKey.generate_key("P-256")
+        key_set = KeySet([rsa_key1, rsa_key2, ec_key])
+
+        algs = JWSRegistry.filter_algorithms(key_set, JWSRegistry.algorithms.keys())
+        names = [alg.name for alg in algs]
+
+        self.assertIn("RS256", names)
+        self.assertIn("ES256", names)
+        self.assertNotIn("ES384", names)
+        self.assertEqual(names.count("RS256"), 1)
+
+    def test_guess_algorithm_with_key_set(self):
+        """guess_algorithm should find the best algorithm across all keys in the KeySet."""
+        rsa_key = RSAKey.generate_key()
+        ec_key = ECKey.generate_key("P-256")
+        key_set = KeySet([rsa_key, ec_key])
+
+        # RS256 comes before ES256 in the recommended list
+        alg = JWSRegistry.guess_algorithm(key_set, JWSRegistry.Strategy.RECOMMENDED)
+        self.assertEqual(alg.name, "RS256")
+
+        # RS512 has the highest algorithm_security (512) among available algorithms
+        alg = JWSRegistry.guess_algorithm(key_set, JWSRegistry.Strategy.SECURITY)
+        self.assertEqual(alg.name, "RS512")