feat: adding import functionality

ape_aws/accounts.py

@@ -1,37 +1,86 @@
 from functools import cached_property
+from json import dumps
+from pathlib import Path
 from typing import Any, Iterator, Optional
 
 from ape.api.accounts import AccountAPI, AccountContainerAPI, TransactionAPI
 from ape.types import AddressType, MessageSignature, SignableMessage, TransactionSignature
+from ape.utils.validators import _validate_account_passphrase
+from eth_account import Account as EthAccount
 from eth_account._utils.legacy_transactions import serializable_unsigned_transaction_from_dict
 from eth_account.messages import _hash_eip191_message, encode_defunct
 from eth_pydantic_types import HexBytes
 from eth_typing import Hash32
-from eth_utils import keccak, to_checksum_address
+from eth_utils import keccak, to_bytes, to_checksum_address
 
 from .client import kms_client
 from .utils import _convert_der_to_rsv
 
 
 class AwsAccountContainer(AccountContainerAPI):
+    loaded_accounts: dict[str, "KmsAccount"] = {}
+
+    def model_post_init(self, __context: Any):
+        print("Initializing AWS KMS Account Container")
+        print([acc.alias for acc in self.accounts])
+
+    @property
+    def _keyfiles(self) -> list[Path]:
+        return [file for file in self.data_folder.glob("*.json")]
+
     @property
     def aliases(self) -> Iterator[str]:
-        return map(lambda x: x.alias, kms_client.raw_aliases)
+        return map(lambda x: x.alias.replace("alias/", ""), kms_client.raw_aliases)
 
     def __len__(self) -> int:
         return len(kms_client.raw_aliases)
 
     @property
     def accounts(self) -> Iterator[AccountAPI]:
+        def _load_account(key_alias, key_id, key_arn) -> AccountAPI:
+            filename = f"{key_alias}.json"
+            keyfile = self.data_folder.joinpath(filename)
+            if filename not in self._keyfiles:
+                self.loaded_accounts[keyfile.stem] = KmsAccount(
+                    key_alias=key_alias,
+                    key_id=key_id,
+                    key_arn=key_arn,
+                )
+                keyfile.write_text(self.loaded_accounts[keyfile.stem].dump_to_json())
+            return self.loaded_accounts[keyfile.stem]
+
         return map(
-            lambda x: KmsAccount(
-                key_alias=x.alias,
+            lambda x: _load_account(
+                key_alias=x.alias.replace("alias/", ""),
                 key_id=x.key_id,
                 key_arn=x.arn,
             ),
             kms_client.raw_aliases,
         )
 
+    def add_private_key(self, alias, passphrase, private_key):
+        kms_account = self.loaded_accounts[alias]
+        _validate_account_passphrase(passphrase)
+        account = EthAccount.from_key(to_bytes(hexstr=private_key))
+        keyfile = self.data_folder.joinpath(f"{alias}.json")
+        account = EthAccount.encrypt(account.key, passphrase)
+        model = kms_account.model_dump()
+        model["address"] = kms_account.address
+        del account["address"]
+        model.update(account)
+        keyfile.write_text(dumps(model, indent=4))
+        print("Key cached successfully")
+        return
+
+    def delete_account(self, alias):
+        alias = alias.replace("alias/", "")
+        keyfile = self.data_folder.joinpath(f"{alias}.json")
+        if keyfile.exists():
+            keyfile.unlink()
+            print(f"Key {alias} deleted successfully")
+        else:
+            print(f"Key {alias} not found")
+
 
 class KmsAccount(AccountAPI):
     key_alias: str
@@ -40,7 +89,7 @@ class KmsAccount(AccountAPI):
 
     @property
     def alias(self) -> str:
-        return self.key_alias.replace("alias/", "")
+        return self.key_alias
 
     @property
     def public_key(self):
@@ -105,3 +154,8 @@ def sign_transaction(self, txn: TransactionAPI, **signer_options) -> Optional[Tr
             )
 
         return txn
+
+    def dump_to_json(self, indent: int = 4):
+        model = self.model_dump()
+        model["address"] = self.address
+        return dumps(model, indent=indent)

ape_aws/client.py

@@ -2,7 +2,11 @@
 from typing import ClassVar
 
 import boto3  # type: ignore[import]
-from pydantic import BaseModel, Field
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import ec, padding
+from eth_account import Account
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 
 class AliasResponse(BaseModel):
@@ -15,10 +19,11 @@ class AliasResponse(BaseModel):
 
 class KeyBaseModel(BaseModel):
     alias: str
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class CreateKeyModel(KeyBaseModel):
-    description: str = Field(alias="Description")
+    description: str | None = Field(default=None, alias="Description")
     policy: str | None = Field(default=None, alias="Policy")
     key_usage: str = Field(default="SIGN_VERIFY", alias="KeyUsage")
     key_spec: str = Field(default="ECC_SECG_P256K1", alias="KeySpec")
@@ -67,10 +72,91 @@ class CreateKey(CreateKeyModel):
     origin: str = Field(default="AWS_KMS", alias="Origin")
 
 
-class ImportKey(CreateKeyModel):
+class ImportKeyRequest(CreateKeyModel):
     origin: str = Field(default="EXTERNAL", alias="Origin")
 
 
+class ImportKey(ImportKeyRequest):
+    key_id: str = Field(default=None, alias="KeyId")
+    public_key: bytes = Field(default=None, alias="PublicKey")
+    private_key: str | bytes | None = Field(default=None, alias="PrivateKey")
+    import_token: bytes = Field(default=None, alias="ImportToken")
+
+    @field_validator("private_key")
+    def validate_private_key(cls, value):
+        if not value:
+            return ec.generate_private_key(ec.SECP256K1(), default_backend())
+        if value.startswith("0x"):
+            value = value[2:]
+        return value
+
+    @property
+    def get_account(self):
+        return Account.privateKeyToAccount(self.private_key)
+
+    @property
+    def ec_private_key(self):
+        loaded_key = self.private_key
+        if isinstance(loaded_key, bytes):
+            loaded_key = ec.derive_private_key(int(self.private_key, 16), ec.SECP256K1())
+        elif isinstance(loaded_key, str):
+            loaded_key = bytes.fromhex(loaded_key[2:])
+            loaded_key = ec.derive_private_key(int(self.private_key, 16), ec.SECP256K1())
+        return loaded_key
+
+    @property
+    def private_key_hex(self):
+        if isinstance(self.private_key, str):
+            return self.private_key
+        elif isinstance(self.private_key, bytes):
+            return self.private_key.hex()
+        return self.private_key.private_numbers().private_value.to_bytes(32, "big").hex()
+
+    @property
+    def private_key_bin(self):
+        """
+        Returns the private key in binary format
+        This is required for the `boto3.client.import_key_material` method
+        """
+        return self.ec_private_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+    @property
+    def private_key_pem(self):
+        """
+        Returns the private key in PEM format for use in outside applications.
+        """
+        return self.ec_private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+    @property
+    def public_key_der(self):
+        return serialization.load_der_public_key(
+            self.public_key,
+            backend=default_backend(),
+        )
+
+    @property
+    def encrypted_private_key(self):
+        if not self.public_key:
+            raise ValueError("Public key not found")
+
+        return self.public_key_der.encrypt(
+            self.private_key_bin,
+            padding.OAEP(
+                mgf=padding.MGF1(hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None,
+            ),
+        )
+
+
 class DeleteKey(KeyBaseModel):
     key_id: str
     days: int = 30
@@ -103,8 +189,9 @@ def sign(self, key_id, msghash):
         )
         return response.get("Signature")
 
-    def create_key(self, key_spec: CreateKey):
+    def create_key(self, key_spec: CreateKey | ImportKeyRequest):
         response = self.client.create_key(**key_spec.to_aws_dict())
+
         key_id = response["KeyMetadata"]["KeyId"]
         self.client.create_alias(
             AliasName=f"alias/{key_spec.alias}",
@@ -131,6 +218,21 @@ def create_key(self, key_spec: CreateKey):
                 )
         return key_id
 
+    def import_key(self, key_spec: ImportKey):
+        return self.client.import_key_material(
+            KeyId=key_spec.key_id,
+            ImportToken=key_spec.import_token,
+            EncryptedKeyMaterial=key_spec.encrypted_private_key,
+            ExpirationModel="KEY_MATERIAL_DOES_NOT_EXPIRE",
+        )
+
+    def get_parameters(self, key_id: str):
+        return self.client.get_parameters_for_import(
+            KeyId=key_id,
+            WrappingAlgorithm="RSAES_OAEP_SHA_256",
+            WrappingKeySpec="RSA_2048",
+        )
+
     def delete_key(self, key_spec: DeleteKey):
         self.client.delete_alias(AliasName=key_spec.alias)
         self.client.schedule_key_deletion(KeyId=key_spec.key_id, PendingWindowInDays=key_spec.days)

ape_aws/kms/_cli.py

@@ -1,7 +1,8 @@
 import click
 from ape.cli import ape_cli_context
 
-from ape_aws.client import CreateKey, DeleteKey, kms_client
+from ape_aws.accounts import AwsAccountContainer, KmsAccount
+from ape_aws.client import CreateKey, DeleteKey, ImportKey, ImportKeyRequest, kms_client
 
 
 @click.group("kms")
@@ -27,14 +28,19 @@ def kms():
     help="Apply key policy to a list of users if applicable, ex. -u ARN1, -u ARN2",
     metavar="list[ARN]",
 )
+@click.option(
+    "-d",
+    "--description",
+    "description",
+    help="The description of the key you intend to create.",
+)
 @click.argument("alias_name")
-@click.argument("description")
 def create_key(
     cli_ctx,
     alias_name: str,
-    description: str,
     administrators: list[str],
     users: list[str],
+    description: str,
 ):
     """
     Create an Ethereum Private Key in AWS KmsAccount
@@ -54,16 +60,90 @@ def create_key(
     cli_ctx.logger.success(f"Key created successfully with ID: {key_id}")
 
 
-# TODO: Add `ape aws kms import`
+@kms.command(name="import")
+@ape_cli_context()
+@click.option(
+    "-p",
+    "--private-key",
+    "private_key",
+    multiple=False,
+    help="The private key to import",
+)
+@click.option(
+    "-a",
+    "--admin",
+    "administrators",
+    multiple=True,
+    help="Apply key policy to a list of administrators if applicable, ex. -a ARN1, -a ARN2",
+    metavar="list[ARN]",
+)
+@click.option(
+    "-u",
+    "--user",
+    "users",
+    multiple=True,
+    help="Apply key policy to a list of users if applicable, ex. -u ARN1, -u ARN2",
+    metavar="list[ARN]",
+)
+@click.option(
+    "-d",
+    "--description",
+    "description",
+    help="The description of the key you intend to create.",
+    metavar="str",
+)
+@click.argument("alias_name")
+def import_key(
+    cli_ctx,
+    alias_name: str,
+    private_key: bytes,
+    administrators: list[str],
+    users: list[str],
+    description: str,
+):
+    def ask_for_passphrase():
+        return click.prompt(
+            "Create Passphrase to encrypt account",
+            hide_input=True,
+            confirmation_prompt=True,
+        )
+
+    passphrase = ask_for_passphrase()
+    key_spec = ImportKeyRequest(
+        alias=alias_name,
+        description=description,  # type: ignore
+        admins=administrators,
+        users=users,
+    )
+    key_id = kms_client.create_key(key_spec)
+    create_key_response = kms_client.get_parameters(key_id)
+    public_key = create_key_response["PublicKey"]
+    import_token = create_key_response["ImportToken"]
+    import_key_spec = ImportKey(
+        **key_spec.model_dump(),
+        key_id=key_id,  # type: ignore
+        public_key=public_key,  # type: ignore
+        private_key=private_key,  # type: ignore
+        import_token=import_token,  # type: ignore
+    )
+    response = kms_client.import_key(import_key_spec)
+    if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
+        cli_ctx.abort("Key failed to import into KMS")
+    cli_ctx.logger.success(f"Key imported successfully with ID: {key_id}")
+    aws_account_container = AwsAccountContainer(name="aws", account_type=KmsAccount)
+    aws_account_container.add_private_key(alias_name, passphrase, import_key_spec.private_key_hex)
+
+
 # TODO: Add `ape aws kms sign-message [message]`
 # TODO: Add `ape aws kms verify-message [message] [hex-signature]`
 
 
 @kms.command(name="delete")
 @ape_cli_context()
 @click.argument("alias_name")
+@click.option("-p", "--purge", is_flag=True, help="Purge the key from the system")
 @click.option("-d", "--days", default=30, help="Number of days until key is deactivated")
-def schedule_delete_key(cli_ctx, alias_name, days):
+def schedule_delete_key(cli_ctx, alias_name, purge, days):
     if "alias" not in alias_name:
         alias_name = f"alias/{alias_name}"
     kms_account = None
@@ -76,4 +156,7 @@ def schedule_delete_key(cli_ctx, alias_name, days):
 
     delete_key_spec = DeleteKey(alias=alias_name, key_id=kms_account.key_id, days=days)
     key_alias = kms_client.delete_key(delete_key_spec)
+    if purge:
+        aws_account_container = AwsAccountContainer(name="aws", account_type=KmsAccount)
+        aws_account_container.delete_account(key_alias)
     cli_ctx.logger.success(f"Key {key_alias} scheduled for deletion in {days} days")