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 `__ + `v2.3.0 `__ ----------------------------------------------------------------------- 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"]