changes to support STS. added in new environment and tags to support KMS and IAM roles. (#70)

Author: twarnock

cloudendure/cloudendure.py

@@ -64,6 +64,10 @@ def __init__(
         self.destination_account: str = self.config.active_config.get(
             "destination_account", ""
         )
+        self.destination_kms: str = self.config.active_config.get("destination_kms", "")
+        self.destination_role: str = self.config.active_config.get(
+            "destination_role", ""
+        )
         self.target_machines: List[str] = self.config.active_config.get(
             "machines", ""
         ).split(",")
@@ -124,6 +128,16 @@ def get_project_id(self, project_name: str = "") -> str:
             return ""
         return project_id
 
+    def _get_role_credentials(self, name: str, role: str) -> Dict[str, Any]:
+        _sts_client: boto_client = boto3.client("sts")
+
+        print(f"Assuming role: {role}")
+        assumed_role: Dict[str, Any] = _sts_client.assume_role(
+            RoleArn=self.destination_role, RoleSessionName=name
+        )
+
+        return assumed_role.get("Credentials")
+
     def check(self) -> bool:
         """Check the status of machines in the provided project."""
         print(
@@ -291,6 +305,8 @@ def update_blueprint(self) -> bool:
                     },
                     {"key": "MigrationWave", "value": self.migration_wave},
                     {"key": "DestinationAccount", "value": self.destination_account},
+                    {"key": "DestinationKMS", "value": self.destination_kms},
+                    {"key": "DestinationRole", "value": self.destination_role},
                 ]
 
                 blueprint["publicIPAction"] = self.config.active_config.get(
@@ -627,25 +643,33 @@ def create_ami(self) -> Dict[str, str]:
             return amis
         return amis
 
-    def copy_image(self, image_id: str, kms_id: str) -> str:
+    def copy_image(self, image_id: str) -> str:
         """Copy a shared image to an account.
 
         Args:
             image_id (str): The AWS AMI to be copied.
-            kms_id (str): The AWS KMS ID to be used for image encryption.
 
         Returns:
             str: The copied AWS AMI ID.
 
         """
-        _ec2_client: boto_client = boto3.client("ec2", AWS_REGION)
+        credentials = self._get_role_credentials("CopyImage", self.destination_role)
+
+        _ec2_client: boto_client = boto3.client(
+            "ec2",
+            region_name=AWS_REGION,
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+        )
 
+        print(f"Copying image {image_id}")
         new_image: Dict[str, Any] = _ec2_client.copy_image(
             SourceImageId=image_id,
             SourceRegion=AWS_REGION,
             Name=f"copied-{image_id}",
             Encrypted=True,
-            KmsKeyId=kms_id,
+            KmsKeyId=self.destination_kms,
         )
 
         return new_image.get("ImageId", "")
@@ -660,8 +684,18 @@ def split_image(self, image_id: str) -> Dict[str, Any]:
             dict: The mapping of AWS EBS block devices.
 
         """
-        print("Loading EC2 resource for region: ", AWS_REGION)
-        _ec2_res = boto3.resource("ec2", AWS_REGION)
+        print(
+            f"Loading EC2 resource for region: {AWS_REGION} using role: {self.destination_role}"
+        )
+        credentials = self._get_role_credentials("SplitImage", self.destination_role)
+
+        _ec2_res: boto_client = boto3.resource(
+            "ec2",
+            region_name=AWS_REGION,
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+        )
 
         # Access the image that needs to be split
         image = _ec2_res.Image(image_id)
@@ -700,6 +734,9 @@ def split_image(self, image_id: str) -> Dict[str, Any]:
                 Tags=[{"Key": f"Drive-{drive}", "Value": json.dumps(drives[drive])}],
             )
 
+        # remove the old image
+        image.deregister()
+
         return root_ami
 
     def gen_terraform(
@@ -726,7 +763,15 @@ def gen_terraform(
             str: The raw Terraform with volume, ENI, and EC2 instance templates.
 
         """
-        _ec2_res = boto3.resource("ec2", AWS_REGION)
+        credentials = self._get_role_credentials("GenTerraform", self.destination_role)
+
+        _ec2_res: boto_client = boto3.resource(
+            "ec2",
+            region_name=AWS_REGION,
+            aws_access_key_id=credentials["AccessKeyId"],
+            aws_secret_access_key=credentials["SecretAccessKey"],
+            aws_session_token=credentials["SessionToken"],
+        )
 
         # Access the image
         image: str = _ec2_res.Image(image_id)

cloudendure/config.py

@@ -33,6 +33,8 @@ class CloudEndureConfig:
         "migration_wave": "0",
         "clone_status": "NOT_STARTED",
         "destination_account": "",
+        "destination_kms": "",
+        "destination_role": "",
         "disk_type": "SSD",
         "public_ip": "DONT_ALLOCATE",
     }

step/iam.tf

@@ -24,6 +24,12 @@ data "aws_iam_policy_document" "lambda-invoke" {
       "*",
     ]
   }
+  # role(s) that the lambdas are allowed to assume roles on for copy, split, and tf generation
+  statement {
+    effect = "Allow"
+    actions = [ "sts:AssumeRole" ]
+    resources = [ "role/arn" ]
+  }
 }
 
 resource "aws_iam_policy" "lambda-invoke" {
@@ -45,7 +51,6 @@ resource "aws_iam_role" "iam_for_lambda" {
 data "aws_iam_policy_document" "iam_for_lambda_assume_role" {
   statement {
     actions = ["sts:AssumeRole"]
-
     principals {
       type        = "Service"
       identifiers = ["lambda.amazonaws.com"]

step/lambdas.tf

@@ -85,12 +85,12 @@ resource "aws_lambda_function" "lambda_split_image" {
   depends_on       = ["data.archive_file.lambdas"]
 }
 
-// resource "aws_lambda_function" "lambda_image_cleanup" {
-//   filename         = "lambdas.zip"
-//   function_name    = "tf-image-cleanup"
-//   role             = "${aws_iam_role.iam_for_lambda.arn}"
-//   handler          = "image_cleanup.lambda_handler"
-//   source_code_hash = "${data.archive_file.lambdas.output_base64sha256}"
-//   runtime          = "python3.7"
-//   depends_on       = ["data.archive_file.lambdas"]
-// }
+resource "aws_lambda_function" "lambda_image_cleanup" {
+  filename         = "lambdas.zip"
+  function_name    = "tf-image-cleanup"
+  role             = "${aws_iam_role.iam_for_lambda.arn}"
+  handler          = "image_cleanup.lambda_handler"
+  source_code_hash = "${data.archive_file.lambdas.output_base64sha256}"
+  runtime          = "python3.7"
+  depends_on       = ["data.archive_file.lambdas"]
+}

step/lambdas/copy_image.py

@@ -26,7 +26,24 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> str:
     migrated_ami_id: str = event["migrated_ami_id"]
     kms_id: str = event["kms_id"]
     region: str = event.get("region", os.environ.get("AWS_REGION"))
-    ec2_client = boto3.client("ec2", region)
+    role: str = event.get("role")
+
+    sts_client = boto3.client("sts")
+
+    print(f"Assuming role: {role}")
+    assumed_role: Dict[str, Any] = sts_client.assume_role(
+        RoleArn=role, RoleSessionName="CopyImageLambda"
+    )
+
+    credentials = assumed_role.get("Credentials")
+
+    ec2_client = boto3.client(
+        "ec2",
+        region_name=region,
+        aws_access_key_id=credentials["AccessKeyId"],
+        aws_secret_access_key=credentials["SecretAccessKey"],
+        aws_session_token=credentials["SessionToken"],
+    )
 
     try:
         new_image: Dict[str, Any] = ec2_client.copy_image(

step/lambdas/find_instance.py

@@ -56,9 +56,12 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> str:
             if tag["Key"] == "DestinationAccount":
                 event_dict["account"] = tag["Value"]
 
-            if tag["Key"] == "KMSId":
+            if tag["Key"] == "DestinationKMS":
                 event_dict["kms_id"] = tag["Value"]
 
+            if tag["Key"] == "DestinationRole":
+                event_dict["role"] = tag["Value"]
+
     except Exception as e:
         print(e)
         event_dict["instance_id"] = "not-found"

step/lambdas/image_cleanup.py

@@ -4,14 +4,13 @@
 from __future__ import annotations
 
 import json
-from typing import Any, Dict
+import os
+from typing import Any, Dict, List
 
 import boto3
 
 print("Loading function image_cleanup")
 
-ec2_client = boto3.client("ec2")
-
 # {
 #   "ami_id": "original AMI",
 #   "kms_id": "KMS GUID used for copy and split AMI",
@@ -25,10 +24,24 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> bool:
     """Handle signaling and entry into the AWS Lambda."""
     print("Received event: " + json.dumps(event, indent=2))
 
-    copy_ami: str = event.get("copy_ami")
-    if not copy_ami:
-        return ""
-
-    ec2_client.deregister_image(ImageId=copy_ami)
+    migrated_ami: str = event.get("migrated_ami_id")
+    if not migrated_ami:
+        return False
+
+    ec2_res = boto3.resource("ec2", os.environ.get("AWS_REGION"))
+    try:
+        # Access the image that needs to be deleted
+        image = ec2_res.Image(migrated_ami)
+
+        # grab device mappings before deregistering
+        devices: List[Any] = image.block_device_mappings
+        image.deregister()
+        for device in devices:
+            if "Ebs" in device:
+                snap = ec2_res.Snapshot(device["Ebs"].get("SnapshotId"))
+                snap.delete()
+    except Exception as e:
+        print(f"Failed.  AMI may not exist.\n{str(e)}")
+        return False
 
     return True

step/lambdas/split_image.py

@@ -25,7 +25,24 @@ def lambda_handler(event: Dict[str, Any], context: Any) -> str:
 
     copy_ami: str = event["copy_ami"]
     region: str = event.get("region", os.environ.get("AWS_REGION"))
-    ec2_res = boto3.resource("ec2", region)
+    role: str = event.get("role")
+
+    sts_client = boto3.client("sts")
+
+    print(f"Assuming role: {role}")
+    assumed_role: Dict[str, Any] = sts_client.assume_role(
+        RoleArn=role, RoleSessionName="SplitImageLambda"
+    )
+
+    credentials = assumed_role.get("Credentials")
+
+    ec2_res = boto3.resource(
+        "ec2",
+        region_name=region,
+        aws_access_key_id=credentials["AccessKeyId"],
+        aws_secret_access_key=credentials["SecretAccessKey"],
+        aws_session_token=credentials["SessionToken"],
+    )
 
     try:
         # Access the image that needs to be split