Add AMI scrub tool

3 files changed, 74 insertions, 2 deletions

diff --git a/.gitignore b/.gitignore
index 98c621a..b9acab3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 /build/
 /.py3/
 /variables.json
+/scrub-old-amis.py
diff --git a/Makefile b/Makefile
index de19494..e7896ca 100644
--- a/Makefile
+++ b/Makefile
@@ -5,7 +5,7 @@ ami: build/convert
 
 build/convert:
         [ -d ".py3" ] || python3 -m venv .py3
-        .py3/bin/pip install pyyaml
+        .py3/bin/pip install pyyaml boto3
 
         [ -d "build" ] || mkdir build
 
@@ -17,6 +17,10 @@ build/convert:
         @echo "json.dump(yaml.load(open(sys.argv[1])), sys.stdout, indent=4, separators=(',', ': '))" >> build/convert
         @chmod +x build/convert
 
+scrub-old-amis.py: scrub-old-amis.py.in
+        sed "s|@PYTHON@|#!`pwd`/.py3/bin/python|" $< > $@
+        chmod +x $@
+
 .PHONY: clean
 clean:
-        rm -rf build .py3
+        rm -rf build .py3 scrub-old-amis.py
diff --git a/scrub-old-amis.py.in b/scrub-old-amis.py.in
new file mode 100644
index 0000000..756f537
--- /dev/null
+++ b/scrub-old-amis.py.in
@@ -0,0 +1,67 @@
+@PYTHON@
+
+import re
+import yaml
+import boto3
+
+
+# All Alpine AMIs should match this regex if they're valid
+AMI_RE = re.compile("^Alpine-(\d+\.\d+)(?:-r(\d+))?-Hardened-EC2")
+
+
+# Load current AMI version from config
+with open("alpine-ami.yaml") as fp:
+    ami_cfg = yaml.load(fp)["variables"]
+    current = (float(ami_cfg["alpine_release"]), int(ami_cfg["ami_release"]))
+
+
+# Fetch all matching AMIs
+amis = []
+
+for region in boto3.session.Session().get_available_regions("ec2"):
+    ec2 = boto3.client("ec2", region_name=region)
+
+    for image in ec2.describe_images(Owners=["self"])["Images"]:
+        match = AMI_RE.match(image["Name"])
+        if not match:
+            continue
+        
+        os_rel, ami_rel = match.groups()
+        amis.append((
+            region, image["ImageId"], float(os_rel),
+            int(ami_rel) if ami_rel else 0))
+
+
+# Determine the set to discard based region and version
+ok_regions = set()
+discards = []
+
+# Cluster candidates by region/version pair, newest in a region first.
+# This should result in the first match for a region always being the newest
+# AMI for that region and all subsequent matches in the region being old.
+# Even so we must keep track of regions with current images on the off-chance
+# that a region only has old images. In that case we want to preserve the old
+# images till we can publish new ones manually so users can still launch
+# Alpine systems without interruption.
+candidates = sorted(amis, key=lambda i: (i[0], (i[1], i[2])), reverse=True)
+
+for ami in candidates:
+    (region, ami), version = ami[:2], ami[2:]
+
+    if version > current:
+        print("{} has AMI '{}' newer than current".format(region, ami))
+        continue
+    elif version == current:
+        ok_regions.add(region)
+        continue
+    elif version < current and region in ok_regions:
+        discards.append((region, ami))
+    else:
+        print("Not discarding old image in {}".format(region))
+        continue
+
+
+# Scrub the old ones
+for region, image in discards:
+    print("Removing image '{}' in {}".format(image, region))
+    boto3.client("ec2", region_name=region).deregister_image(ImageId=image)