Add exact_count and count_tag to the ec2 module.

cloud/ec2

@@ -376,6 +376,7 @@ local_action:
 
 import sys
 import time
+from ast import literal_eval
 
 try:
     import boto.ec2
@@ -385,6 +386,68 @@ except ImportError:
     print "failed=True msg='boto required for this module'"
     sys.exit(1)
 
+def find_running_instances_by_count_tag(module, ec2, count_tag):
+
+    # get reservations for instances that match tag(s) and are running
+    reservations = get_reservations(module, ec2, tags=count_tag, state="running")
+
+    instances = []
+    for res in reservations:
+        if hasattr(res, 'instances'):
+            for inst in res.instances:
+                instances.append(inst) 
+                
+    return reservations, instances
+
+
+def _set_none_to_blank(dictionary):
+    result = dictionary
+    for k in result.iterkeys():
+        if type(result[k]) == dict:
+            result[k] = _set_non_to_blank(result[k])
+        elif not result[k]:
+            result[k] = ""
+    return result                        
+
+
+def get_reservations(module, ec2, tags=None, state=None):
+
+    # TODO: filters do not work with tags that have underscores
+    filters = dict()
+
+    if tags is not None:
+
+        if type(tags) is str:
+            try:
+                tags = literal_eval(tags)
+            except:
+                pass
+
+        # if string, we only care that a tag of that name exists
+        if type(tags) is str:
+            filters.update({"tag-key": tags})
+
+        # if list, append each item to filters
+        if type(tags) is list:
+            for x in tags:
+                if type(x) is dict:
+                    x = _set_none_to_blank(x)
+                    filters.update(dict(("tag:"+tn, tv) for (tn,tv) in x.iteritems()))
+                else:
+                    filters.update({"tag-key": x})
+
+        # if dict, add the key and value to the filter
+        if type(tags) is dict:
+            tags = _set_none_to_blank(tags)
+            filters.update(dict(("tag:"+tn, tv) for (tn,tv) in tags.iteritems()))
+
+    if state:
+        # http://stackoverflow.com/questions/437511/what-are-the-valid-instancestates-for-the-amazon-ec2-api
+        filters.update({'instance-state-name': state})
+
+    results = ec2.get_all_instances(filters=filters)
+
+    return results
 
 def get_instance_info(inst):
     """
@@ -473,7 +536,45 @@ def create_block_device(module, ec2, volume):
                            delete_on_termination=volume.get('delete_on_termination', False),
                            iops=volume.get('iops'))
 
-def create_instances(module, ec2):
+
+def enforce_count(module, ec2):
+
+    exact_count = module.params.get('exact_count')
+    count_tag = module.params.get('count_tag')
+
+    reservations, instances = find_running_instances_by_count_tag(module, ec2, count_tag)
+
+    changed = None
+    checkmode = False
+    instance_dict_array = None
+    changed_instance_ids = None
+
+    if len(instances) == exact_count:
+        changed = False
+    elif len(instances) < exact_count:
+        changed = True
+        to_create = exact_count - len(instances)
+        if not checkmode:
+            (instance_dict_array, changed_instance_ids, changed) \
+                = create_instances(module, ec2, override_count=to_create)
+    elif len(instances) > exact_count:
+        changed = True
+        to_remove = len(instances) - exact_count
+        if not checkmode:
+            all_instance_ids = sorted([ x.id for x in instances ])
+            remove_ids = all_instance_ids[0:to_remove]
+            (changed, instance_dict_array, changed_instance_ids) \
+                = terminate_instances(module, ec2, remove_ids)
+            terminated_list = []
+            for inst in instance_dict_array:
+                inst['state'] = "terminated"
+                terminated_list.append(inst)
+            instance_dict_array = terminated_list            
+
+    return (instance_dict_array, changed_instance_ids, changed)
+    
+        
+def create_instances(module, ec2, override_count=None):
     """
     Creates new instances
 
@@ -492,7 +593,10 @@ def create_instances(module, ec2):
     zone = module.params.get('zone')
     instance_type = module.params.get('instance_type')
     image = module.params.get('image')
-    count = module.params.get('count')
+    if override_count:
+        count = override_count
+    else:
+        count = module.params.get('count')
     monitoring = module.params.get('monitoring')
     kernel = module.params.get('kernel')
     ramdisk = module.params.get('ramdisk')
@@ -506,6 +610,8 @@ def create_instances(module, ec2):
     private_ip = module.params.get('private_ip')
     instance_profile_name = module.params.get('instance_profile_name')
     volumes = module.params.get('volumes')
+    exact_count = module.params.get('exact_count')
+    count_tag = module.params.get('count_tag')
 
     # group_id and group_name are exclusive of each other
     if group_id and group_name:
@@ -832,6 +938,8 @@ def main():
             instance_profile_name = dict(),
             instance_ids = dict(type='list'),
             state = dict(default='present'),
+            exact_count = dict(type='int'),
+            count_tag = dict(),
             volumes = dict(type='list'),
         )
     )
@@ -857,7 +965,11 @@ def main():
         # Changed is always set to true when provisioning new instances
         if not module.params.get('image'):
             module.fail_json(msg='image parameter is required for new instance')
-        (instance_dict_array, new_instance_ids, changed) = create_instances(module, ec2)
+    
+        if module.params.get('exact_count'):
+            (instance_dict_array, new_instance_ids, changed) = enforce_count(module, ec2)
+        else:            
+            (instance_dict_array, new_instance_ids, changed) = create_instances(module, ec2)
 
     module.exit_json(changed=changed, instance_ids=new_instance_ids, instances=instance_dict_array)