api_jwk: Add PyJWKSet.__getitem__ (#725)

* api_jwk: Add PyJWKSet.__getitem__

Closes #724.

* CHANGELOG: record changes

3 files changed, 31 insertions, 0 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 498a544..06d281c 100644
--- a/CHANGELOG.rst
+++ b/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>`__
 -----------------------------------------------------------------------
 
diff --git a/jwt/api_jwk.py b/jwt/api_jwk.py
index a0f6364..f58bae2 100644
--- a/jwt/api_jwk.py
+++ b/jwt/api_jwk.py
@@ -95,3 +95,9 @@ class PyJWKSet:
     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)
diff --git a/tests/test_api_jwk.py b/tests/test_api_jwk.py
index 53c0547..6f6cb89 100644
--- a/tests/test_api_jwk.py
+++ b/tests/test_api_jwk.py
@@ -252,3 +252,26 @@ class TestPyJWKSet:
         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"]