Repository: jpadilla/pyjwt

Issue:
New API: Retrieving a key by KID from a `PyJWKSet`?
First of all, thanks for creating and maintaining this library!

The current `PyJWKSet` API is pretty thin: it wraps a list of `PyJWK` and provides basic deserialization, but nothing else.

I'd like to propose adding `PyJWKSet.__getitem__`, which would allow a user to index the `PyJWKSet` by key ID and return the corresponding `PyJWK` (or `KeyError` if not present). This saves users from having to write their own search over the `PyJWKSet`, and doesn't significantly complicate the current API.

With `__getitem__`, the use pattern would be:

```python
jwks = PyJWKSet(keys)
jwk = jwks["somelonghexkeyid"]
```

Would this functionality be welcome? If so, I am happy to contribute it. 

Discussion:
woodruffw: Gentle ping on this. Anything I can do to help with review?

Patch:
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)


Test patch:
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"]