feat: adding import functionality

ape_aws/client.py

@@ -1,8 +1,13 @@
+from cryptography.hazmat.primitives.asymmetric import ec, padding
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.backends import default_backend
+from eth_account import Account
+
 from datetime import datetime
 from typing import ClassVar
 
 import boto3  # type: ignore[import]
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict, field_validator
 
 
 class AliasResponse(BaseModel):
@@ -15,6 +20,7 @@ class AliasResponse(BaseModel):
 
 class KeyBaseModel(BaseModel):
     alias: str
+    model_config = ConfigDict(populate_by_name=True)
 
 
 class CreateKeyModel(KeyBaseModel):
@@ -67,10 +73,76 @@ 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: 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 isinstance(value, bytes):
+            return ec.generate_private_key(
+                ec.SECP256K1(),
+                default_backend()
+            )
+        if value.startswith('0x'):
+            return value[2:]
+        return value
+
+    @property
+    def get_account(self):
+        return Account.privateKeyToAccount(self.private_key)
+
+    @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.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.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 +175,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 | ImportKey):
         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 +204,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,13 @@
 import click
-from ape.cli import ape_cli_context
 
-from ape_aws.client import CreateKey, DeleteKey, kms_client
+from ape.cli import ape_cli_context
+from ape_aws.client import (
+    CreateKey,
+    DeleteKey,
+    ImportKeyRequest,
+    ImportKey,
+    kms_client,
+)
 
 
 @click.group("kms")
@@ -54,7 +60,62 @@ 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.argument("alias_name")
+@click.argument("description")
+def import_key(
+    cli_ctx,
+    alias_name: str,
+    description: str,
+    private_key: bytes,
+    administrators: list[str],
+    users: list[str],
+):
+    key_spec = ImportKeyRequest(
+        alias=alias_name,
+        description=description,
+        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,
+        public_key=public_key,
+        private_key=private_key,
+        import_token=import_token,
+    )
+    key_id = kms_client.import_key(import_key_spec)
+    cli_ctx.logger.success(f"Key imported successfully with ID: {key_id}")
+
+
 # TODO: Add `ape aws kms sign-message [message]`
 # TODO: Add `ape aws kms verify-message [message] [hex-signature]`