get_copy_status now uses role assumption (#89)

Author: twarnock

step/lambdas/get_copy_status.py

@@ -29,18 +29,39 @@ 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"))
     instance_name: str = event.get("name", "")
-    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="GetCopyStatusLambda"
+    )
+
+    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:
         ami_state: Dict[str, Any] = ec2_client.describe_images(ImageIds=[copy_ami])
     except Exception as e:
         print(e)
         return "error"
 
-    state = ami_state["Images"][0]["State"]
-    if state == "available":
+    images = ami_state.get("Images")
+    if not images:
+        return "no-image-yet"
+
+    state = images[0].get("State")
+    if state and state == "available":
         MigrationStateHandler().update_state(
             state="IMAGE_COPIED", machine_name=instance_name
         )
 
-    return ami_state["Images"][0]["State"]
+    return state