From de3e55f9c63e87b6a68b09322a83babebc39304f Mon Sep 17 00:00:00 2001
From: Blake Johnson <blakeejohnson39@gmail.com>
Date: Thu, 25 Apr 2024 15:52:59 -0500
Subject: [PATCH] feat: adding key import, not finished

---
 ape_aws/_cli.py         |  4 +-
 ape_aws/accounts.py     |  2 +-
 ape_aws/kms/__init__.py |  0
 ape_aws/kms/_cli.py     | 96 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 ape_aws/kms/__init__.py
 create mode 100644 ape_aws/kms/_cli.py

diff --git a/ape_aws/_cli.py b/ape_aws/_cli.py
index 65b8585..7d75980 100644
--- a/ape_aws/_cli.py
+++ b/ape_aws/_cli.py
@@ -1,9 +1,11 @@
 import click
 
+from ape_aws.kms._cli import kms
+
 
 @click.group()
 def cli():
     """Ape AWS CLI commands"""
 
 
-cli.add_command()
+cli.add_command(kms)
diff --git a/ape_aws/accounts.py b/ape_aws/accounts.py
index e780a61..489906d 100644
--- a/ape_aws/accounts.py
+++ b/ape_aws/accounts.py
@@ -12,7 +12,7 @@
 from ape.types import AddressType, MessageSignature, SignableMessage
 from ape.utils import cached_property
 
-from .utils import SECP256_K1_N, AliasResponse, transaction
+from .utils import SECP256_K1_N, AliasResponse
 
 
 class AwsAccountContainer(AccountContainerAPI):
diff --git a/ape_aws/kms/__init__.py b/ape_aws/kms/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/ape_aws/kms/_cli.py b/ape_aws/kms/_cli.py
new file mode 100644
index 0000000..c5ec155
--- /dev/null
+++ b/ape_aws/kms/_cli.py
@@ -0,0 +1,96 @@
+import boto3
+import click
+
+from ape.cli import ape_cli_context
+
+from ape_aws.accounts import KmsAccount, AwsAccountContainer
+
+
+@click.group("kms")
+def kms():
+    """Manage AWS KMS keys"""
+
+
+@kms.command()
+@ape_cli_context
+@click.argument("alias")
+def account(alias):
+    """
+    Import an existing Ethereum Private Key into AWS KMS
+    """
+    aws_accounts = AwsAccountContainer(data_folder='./', account_type=KmsAccount)
+    kms_account = None
+    for account in list(aws_accounts.accounts):
+        if account.key_alias == alias:
+            kms_account = account
+
+    if not kms_account:
+        raise ValueError(f"No KMS Key with alias name {alias}")
+
+    return kms_account
+
+
+@kms.command()
+@ape_cli_context
+@click.argument("alias_name")
+@click.argument("description")
+@click.argument("tag_key")
+@click.argument("tag_value")
+@click.argument("administrators")
+@click.argument("users")
+def import_key(alias_name, description, tag_key, tag_value, administrators, users):
+    """
+    Import an existing Ethereum Private Key into AWS KMS
+    """
+    aws_account = AwsAccountContainer(data_folder='./', account_type=KmsAccount)
+    response = aws_account.kms_client.create_key(
+        KeyUsage='SIGN_VERIFY',
+        KeySpec='ECC_SECG_P256K1',
+        Origin='External',
+        MultiRegion=False,
+    )
+    key_id = response['KeyMetadata']['KeyId']
+    aws_account.kms_client.create_alias(
+        AliasName=f'alias/{alias_name}',
+        TargetKeyId=key_id,
+    )
+    aws_account.kms_client.tag_resource(
+        KeyId=key_id,
+        Tags=[{'TagKey': tag_key, 'TagValue': tag_value}],
+    )
+    # Note: Get the ARN from AWS
+    for arn in administrators:
+        aws_account.kms_client.put_key_policy(
+            KeyId=key_id,
+            PolicyName='default',
+            Policy='''{
+                "Version": "2012-10-17",
+                "Id": "key-default-1",
+                "Statement": [{
+                    "Sid": "Enable IAM User Permissions",
+                    "Effect": "Allow",
+                    "Principal": {"AWS": "%s"},
+                    "Action": "kms:*",
+                    "Resource": "*"
+                }]
+            }''' % arn
+        )
+    # Note: get ARN from AWS
+    for arn in users:
+        aws_account.kms_client.put_key_policy(
+            KeyId=key_id,
+            PolicyName='default',
+            Policy='''{
+                "Version": "2012-10-17",
+                "Id": "key-default-1",
+                "Statement": [{
+                    "Sid": "Allow use of the key",
+                    "Effect": "Allow",
+                    "Principal": {"AWS": "%s"},
+                    "Action": ["kms:Sign", "kms:Verify"],
+                    "Resource": "*"
+                }]
+            }''' % arn
+        )
+
+    print("Key created successfully with ID: ", key_id)