api_jwk: Add PyJWKSet.__getitem__ (#725)

* api_jwk: Add PyJWKSet.__getitem__

Closes #724.

* CHANGELOG: record changes

CHANGELOG.rst

@@ -16,6 +16,8 @@ Fixed
 Added
 ~~~~~
 
+- Add ``PyJWKSet.__getitem__`` for indexing keysets by key ID `#725 <https://github.com/jpadilla/pyjwt/pull/725>`__
+
 `v2.3.0 <https://github.com/jpadilla/pyjwt/compare/2.2.0...2.3.0>`__
 -----------------------------------------------------------------------
 

jwt/api_jwk.py

@@ -95,3 +95,9 @@ def from_dict(obj):
     def from_json(data):
         obj = json.loads(data)
         return PyJWKSet.from_dict(obj)
+
+    def __getitem__(self, kid):
+        for key in self.keys:
+            if key.key_id == kid:
+                return key
+        raise KeyError("keyset has no key for kid: %s" % kid)

tests/test_api_jwk.py

@@ -252,3 +252,26 @@ def test_should_load_keys_from_jwk_data_json_string(self):
         assert jwk.key_type == "RSA"
         assert jwk.key_id == "keyid-abc123"
         assert jwk.public_key_use == "sig"
+
+    @crypto_required
+    def test_keyset_should_index_by_kid(self):
+        algo = RSAAlgorithm(RSAAlgorithm.SHA256)
+
+        with open(key_path("jwk_rsa_pub.json")) as keyfile:
+            pub_key = algo.from_jwk(keyfile.read())
+
+        key_data_str = algo.to_jwk(pub_key)
+        key_data = json.loads(key_data_str)
+
+        # TODO Should `to_jwk` set these?
+        key_data["alg"] = "RS256"
+        key_data["use"] = "sig"
+        key_data["kid"] = "keyid-abc123"
+
+        jwk_set = PyJWKSet.from_dict({"keys": [key_data]})
+
+        jwk = jwk_set.keys[0]
+        assert jwk == jwk_set["keyid-abc123"]
+
+        with pytest.raises(KeyError):
+            jwk_set["this-kid-does-not-exist"]