From ae8ac14ed738881274057f5b7d4361ce475fb362 Mon Sep 17 00:00:00 2001
From: William Woodruff <william@trailofbits.com>
Date: Fri, 21 Jan 2022 10:44:54 -0500
Subject: [PATCH 1/2] api_jwk: Add PyJWKSet.__getitem__

Closes #724.
---
 jwt/api_jwk.py        |  6 ++++++
 tests/test_api_jwk.py | 23 +++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/jwt/api_jwk.py b/jwt/api_jwk.py
index a0f6364d..f58bae28 100644
--- a/jwt/api_jwk.py
+++ b/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)
diff --git a/tests/test_api_jwk.py b/tests/test_api_jwk.py
index 53c05473..6f6cb899 100644
--- a/tests/test_api_jwk.py
+++ b/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"]

From c593430aa7d4f4878142ddc9f5adb59fee829797 Mon Sep 17 00:00:00 2001
From: William Woodruff <william@trailofbits.com>
Date: Fri, 21 Jan 2022 10:52:13 -0500
Subject: [PATCH 2/2] CHANGELOG: record changes

---
 CHANGELOG.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 498a544a..06d281c1 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>`__
 -----------------------------------------------------------------------