From 5898f9e34869b654a48b05440fffe91e9dccd9b8 Mon Sep 17 00:00:00 2001
From: Arian van Putten <arian.vanputten@gmail.com>
Date: Tue, 21 May 2024 22:47:57 +0200
Subject: [PATCH] Some fixes

---
 upload-ami/src/upload_ami/delete_images.py | 15 +++---
 upload-ami/src/upload_ami/nuke.py          | 55 ++++++++++++++--------
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/upload-ami/src/upload_ami/delete_images.py b/upload-ami/src/upload_ami/delete_images.py
index c9ad01f..843e664 100644
--- a/upload-ami/src/upload_ami/delete_images.py
+++ b/upload-ami/src/upload_ami/delete_images.py
@@ -6,7 +6,7 @@
 logger = logging.getLogger(__name__)
 
 
-def delete_images_by_name(ec2: EC2Client, image_name: str) -> None:
+def delete_images_by_name(ec2: EC2Client, image_name: str, dry_run: bool) -> None:
     """
     Delete an image by its name.
 
@@ -19,7 +19,6 @@ def delete_images_by_name(ec2: EC2Client, image_name: str) -> None:
         OwnerIds=["self"], Filters=[{"Name": "tag:Name", "Values": [image_name]}]
     )
     logger.info(f"Deleting {len(snapshots['Snapshots'])} snapshots")
-    input("Press Enter to continue")
 
     for snapshot in snapshots["Snapshots"]:
         assert "SnapshotId" in snapshot
@@ -37,9 +36,9 @@ def delete_images_by_name(ec2: EC2Client, image_name: str) -> None:
         for image in images["Images"]:
             assert "ImageId" in image
             logger.info(f"Deregistering {image['ImageId']}")
-            ec2.deregister_image(ImageId=image["ImageId"])
+            ec2.deregister_image(ImageId=image["ImageId"], DryRun=dry_run)
         logger.info(f"Deleting {snapshot['SnapshotId']}")
-        ec2.delete_snapshot(SnapshotId=snapshot["SnapshotId"])
+        ec2.delete_snapshot(SnapshotId=snapshot["SnapshotId"], DryRun=dry_run)
 
 
 def main() -> None:
@@ -54,11 +53,15 @@ def main() -> None:
         "--all-regions",
         action="store_true",
     )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+    )
     logging.basicConfig(level=logging.INFO)
     ec2: EC2Client = boto3.client("ec2")
 
     args = parser.parse_args()
-    delete_images_by_name(ec2, args.image_name)
+    delete_images_by_name(ec2, args.image_name, args.dry_run)
     if args.all_regions:
         regions = ec2.describe_regions()["Regions"]
         for region in regions:
@@ -67,7 +70,7 @@ def main() -> None:
             logger.info(
                 f"Deleting image by name {args.image_name} in {region['RegionName']}"
             )
-            delete_images_by_name(ec2r, args.image_name)
+            delete_images_by_name(ec2r, args.image_name, args.dry_run)
 
 
 if __name__ == "__main__":
diff --git a/upload-ami/src/upload_ami/nuke.py b/upload-ami/src/upload_ami/nuke.py
index 2fd53ea..1ba5829 100644
--- a/upload-ami/src/upload_ami/nuke.py
+++ b/upload-ami/src/upload_ami/nuke.py
@@ -1,9 +1,27 @@
 import logging
 import boto3
 from mypy_boto3_ec2 import EC2Client
+import argparse
+import botocore.exceptions
 
 
 def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--image-name",
+        type=str,
+        required=True,
+        help="Name of the image to delete. Can be a filter.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--older-than",
+        type=str,
+    )
+    args = parser.parse_args()
     logging.basicConfig(level=logging.INFO)
     ec2: EC2Client = boto3.client("ec2", region_name="us-east-1")
 
@@ -13,26 +31,23 @@ def main() -> None:
         assert "RegionName" in region
         ec2r = boto3.client("ec2", region_name=region["RegionName"])
         logging.info(f"Nuking {region['RegionName']}")
-        snapshots = ec2r.describe_snapshots(OwnerIds=["self"])
-        for snapshot in snapshots["Snapshots"]:
-
-            assert "SnapshotId" in snapshot
-            images = ec2r.describe_images(
-                Owners=["self"],
-                Filters=[
-                    {
-                        "Name": "block-device-mapping.snapshot-id",
-                        "Values": [snapshot["SnapshotId"]],
-                    }
-                ],
-            )
-            for image in images["Images"]:
-                assert "ImageId" in image
-                logging.info(f"Deregistering {image['ImageId']}")
-                ec2r.deregister_image(ImageId=image["ImageId"])
-
-            logging.info(f"Deleting {snapshot['SnapshotId']}")
-            ec2r.delete_snapshot(SnapshotId=snapshot["SnapshotId"])
+        images = ec2r.describe_images(
+            Owners=["self"], Filters=[{"Name": "name", "Values": [args.image_name]}]
+        )
+        for image in images["Images"]:
+            snapshot_id = image["BlockDeviceMappings"][0]["Ebs"]["SnapshotId"]
+            logging.info(f"Deregistering {image['ImageId']}")
+            try:
+                ec2r.deregister_image(ImageId=image["ImageId"], DryRun=args.dry_run)
+            except botocore.exceptions.ClientError as e:
+                if "DryRunOperation" not in str(e):
+                    raise
+            logging.info(f"Deleting {snapshot_id}")
+            try:
+                ec2r.delete_snapshot(SnapshotId=snapshot_id, DryRun=args.dry_run)
+            except botocore.exceptions.ClientError as e:
+                if "DryRunOperation" not in str(e):
+                    raise
 
 
 if __name__ == "__main__":