fix(KMS): handle deleted key better (#28)

Author: fubuloubu

ape_aws/accounts.py

@@ -25,13 +25,15 @@ def __init__(self, *args, **kwargs):
     @property
     def keys(self) -> dict[str, KmsKey]:  # type: ignore[syntax]
         try:
-            return super(AwsClient, self).keys
+            keys = super(AwsClient, self).keys
 
         except AwsAccessError as e:
             # NOTE: Do not raise here, instead just log warning (prevent issues w/ Ape API)
             logger.warning(str(e))
             return {}
 
+        return {alias: key for alias, key in keys.items() if key.enabled}
+
     @property
     def aliases(self) -> Iterator[str]:
         yield from iter(self.keys)

ape_aws/iam/client.py

@@ -98,7 +98,7 @@ class IamClient(Session):
     KEY_ACCESS_POLICY = dict(
         Sid="ApeAWSv1",
         Effect="Allow",
-        Action=["kms:ListAliases", "kms:Sign", "kms:Verify", "kms:GetPublicKey"],
+        Action=["kms:ListAliases", "kms:Sign", "kms:Verify", "kms:GetPublicKey", "kms:DescribeKey"],
         Resource="*",
     )
 
@@ -115,7 +115,7 @@ def users(self) -> dict[str, IamUser]:
 
         except BotoCoreError as e:
             # NOTE: Handle here since `.users` is the main access point for the external API
-            raise AwsAccessError(e)
+            raise AwsAccessError(e) from e
 
         users = map(IamUser.model_validate, response.get("Users", []))
 

ape_aws/kms/client.py

@@ -3,7 +3,7 @@
 from itertools import chain
 from typing import TYPE_CHECKING, ClassVar
 
-from ape.types import AddressType
+from ape.types import AddressType, HexBytes
 from botocore.exceptions import BotoCoreError  # type: ignore[import-untyped]
 from pydantic import BaseModel, Field, SecretStr, field_validator
 
@@ -76,6 +76,14 @@ class KmsKey(BaseModel):
     def prune_alias_prepend(cls, value: str):
         return value.replace("alias/ape-aws/v1/", "")
 
+    @property
+    def metadata(self) -> dict:
+        return self.kms_client.describe_key(KeyId=self.id)["KeyMetadata"]
+
+    @property
+    def enabled(self) -> bool:
+        return self.metadata.get("Enabled", False)
+
     @property
     def alias(self) -> str:
         if self.cached_alias is not None:
@@ -96,8 +104,12 @@ def set_alias(self, alias: str):
         self.cached_alias = alias
 
     @property
-    def public_key(self):
-        return self.kms_client.get_public_key(KeyId=self.id)["PublicKey"]
+    def public_key(self) -> HexBytes:
+        try:
+            return HexBytes(self.kms_client.get_public_key(KeyId=self.id)["PublicKey"][-64:])
+        except (self.kms_client.exceptions.KMSInvalidStateException, BotoCoreError) as e:
+            # NOTE: Handle here since `.keys` is the main access point for the external API
+            raise AwsAccessError(e) from e
 
     @property
     def address(self) -> AddressType:
@@ -106,7 +118,7 @@ def address(self) -> AddressType:
 
         from eth_utils import keccak, to_checksum_address
 
-        return AddressType(to_checksum_address(keccak(self.public_key[-64:])[-20:].hex().lower()))
+        return AddressType(to_checksum_address(keccak(self.public_key)[-20:]))
 
     def add_tags(self, tags: list[dict[str, str]]):
         self.kms_client.tag_resource(KeyId=self.id, Tags=tags)
@@ -189,7 +201,7 @@ def keys(self) -> dict[str, KmsKey]:
             key_data = filter(lambda k: k["AliasName"].startswith("alias/ape-aws/"), chain(*pages))
         except BotoCoreError as e:
             # NOTE: Handle here since `.keys` is the main access point for the external API
-            raise AwsAccessError(e)
+            raise AwsAccessError(e) from e
 
         keys = map(KmsKey.model_validate, key_data)
         return {key.alias: key for key in keys}