From a5cd6e077d57ccace73af65404c7acc38d8bfea9 Mon Sep 17 00:00:00 2001
From: Jonh Wendell <jonh.wendell@gmail.com>
Date: Mon, 6 Apr 2015 14:23:36 -0300
Subject: [PATCH] ec2_launch: Allow to deregister the AMI after launching image

If all you want is to test an image or product and to achieve this
you need to generate several images a day, you will end up with
lots of AMI's and snapshots that have no use in the end of the day.

This commit adds the new boolean manifest option "deregister_ami"
that, if True, deletes the recently created AMI and snapshot. So, the
final result will be only the running instance, nothing else.
---
 bootstrapvz/plugins/ec2_launch/__init__.py    |  2 ++
 .../plugins/ec2_launch/manifest-schema.yml    |  2 ++
 bootstrapvz/plugins/ec2_launch/tasks.py       | 21 +++++++++++++++++++
 3 files changed, 25 insertions(+)

diff --git a/bootstrapvz/plugins/ec2_launch/__init__.py b/bootstrapvz/plugins/ec2_launch/__init__.py
index c1e8caf..69c29c7 100644
--- a/bootstrapvz/plugins/ec2_launch/__init__.py
+++ b/bootstrapvz/plugins/ec2_launch/__init__.py
@@ -9,3 +9,5 @@ def resolve_tasks(taskset, manifest):
     taskset.add(tasks.LaunchEC2Instance)
     if 'print_public_ip' in manifest.plugins['ec2_launch']:
         taskset.add(tasks.PrintPublicIPAddress)
+    if manifest.plugins['ec2_launch'].get('deregister_ami', False):
+        taskset.add(tasks.DeregisterAMI)
diff --git a/bootstrapvz/plugins/ec2_launch/manifest-schema.yml b/bootstrapvz/plugins/ec2_launch/manifest-schema.yml
index faba1e8..9d7992e 100644
--- a/bootstrapvz/plugins/ec2_launch/manifest-schema.yml
+++ b/bootstrapvz/plugins/ec2_launch/manifest-schema.yml
@@ -16,3 +16,5 @@ properties:
           instance_type: {type: string}
           print_public_ip: {type: string}
           tags: {type: object}
+          deregister_ami: {type: boolean}
+        additionalProperties: false
diff --git a/bootstrapvz/plugins/ec2_launch/tasks.py b/bootstrapvz/plugins/ec2_launch/tasks.py
index c597254..5d4abc5 100644
--- a/bootstrapvz/plugins/ec2_launch/tasks.py
+++ b/bootstrapvz/plugins/ec2_launch/tasks.py
@@ -62,3 +62,24 @@ class PrintPublicIPAddress(Task):
             f.write('')
 
         f.close()
+
+
+class DeregisterAMI(Task):
+    description = 'Deregistering AMI'
+    phase = phases.image_registration
+    predecessors = [LaunchEC2Instance]
+
+    @classmethod
+    def run(cls, info):
+        ec2 = info._ec2
+        logger = logging.getLogger(__name__)
+
+        def instance_running():
+            ec2['instance'].update()
+            return ec2['instance'].state == 'running'
+
+        if waituntil(instance_running, timeout=120, interval=5):
+            info._ec2['connection'].deregister_image(info._ec2['image'])
+            info._ec2['snapshot'].delete()
+        else:
+            logger.error('Timeout while booting instance')