initial import

[awsible] / inventory / asg-inventory.py

diff --git a/inventory/asg-inventory.py b/inventory/asg-inventory.py
new file mode 100755 (executable)
index 0000000..74e21b0
--- /dev/null
+++ b/inventory/asg-inventory.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python
+'''\
+Generate a JSON object containing the names of all the AWS Autoscaling
+Groups in an account and the IPs of the Instances within them, suitable
+for use as an Ansible inventory.
+'''
+
+import argparse
+import boto3
+import json
+import sys
+from multiprocessing.dummy import Pool as ThreadPool
+from functools import partial
+
+
+DEFAULT_REGIONS = ['us-east-1', 'us-west-2']
+HOSTVARS = {}
+
+
+def allASGInstances(asgc):
+    'Return a tuple of a dict of each ASG name listing the instance IDs within it, and a list of all instance IDs.'
+    asgs = {}
+    instanceIds = []
+    args = {}
+    while True:
+        response = asgc.describe_auto_scaling_groups(**args)
+        for asg in response['AutoScalingGroups']:
+            asgs[asg['AutoScalingGroupName']] = [i['InstanceId'] for i in asg['Instances']]
+            instanceIds += asgs[asg['AutoScalingGroupName']]
+        if 'NextToken' not in response:
+            break
+        args = {'NextToken': response['NextToken']}
+    return (asgs, instanceIds)
+
+
+def allInstanceIPs(ec2c, InstanceIds=None, publicIPs=False):
+    'Return a dict of each Instance ID with its addresses.'
+    instances = {}
+    args = {}
+    IPType = 'PublicIpAddress' if publicIPs else 'PrivateIpAddress'
+    if InstanceIds is not None:
+        args['InstanceIds'] = InstanceIds
+    while True:
+        response = ec2c.describe_instances(**args)
+        for resv in response['Reservations']:
+            for inst in resv['Instances']:
+                if IPType in inst:
+                    instances[inst['InstanceId']] = inst[IPType]
+        if 'NextToken' not in response:
+            break
+        args = {'NextToken': response['NextToken']}
+    return instances
+
+
+def regionInventory(sessionArgs, publicIPs=False):
+    'Return dict results for one region.'
+    session = boto3.session.Session(**sessionArgs)
+    asgc = session.client('autoscaling')
+    ec2c = session.client('ec2')
+
+    # get dict of ASG names and associated Instance Ids, plus list of all Instance Ids referenced by ASGs
+    (ASGs, AllInstanceIds) = allASGInstances(asgc)
+
+    # get list of instance IPs for all instance Ids used by ASGs
+    AllInstanceIPs = allInstanceIPs(ec2c, InstanceIds=AllInstanceIds, publicIPs=publicIPs)
+
+    # a group for every Instance Id
+    inventory = {iid:[AllInstanceIPs[iid]] for iid in AllInstanceIPs}
+
+    # add ASG dict, replacing ASG Instance Id with instance IP
+    inventory.update({asg:[AllInstanceIPs[iid] for iid in ASGs[asg]] for asg in ASGs})
+
+    return inventory
+
+
+def mergeDictOfLists(a, b):
+    'There is likely a better way of doing this, but right now I have a headache.'
+    for key in b:
+        if key in a:
+            a[key] += b[key]
+        else:
+            a[key] = b[key]
+    return a
+
+
+parser = argparse.ArgumentParser(description='dynamic Ansible inventory from AWS Autoscaling Groups')
+parser.add_argument('--public', action='store_true', help='inventory public IPs (default: private IPs)')
+parser.add_argument('--profile', metavar='PROFILE', dest='profile_name', help='AWS Profile (default: current IAM Role)')
+parser.add_argument('--regions', metavar='REGION', nargs='+', default=DEFAULT_REGIONS, help='AWS Regions (default: %(default)s)')
+parser.add_argument('--list', action='store_true')
+parser.add_argument('--host', nargs=1)
+args = parser.parse_args()
+
+if args.host:
+    print(json.dumps(HOSTVARS))
+    sys.exit()
+
+# create sessionArgs for each region
+regionArgs = [{'region_name': region} for region in args.regions]
+if args.profile_name:
+    for arg in regionArgs:
+        arg.update({'profile_name': args.profile_name})
+
+# pin the non-variant option
+invf = partial(regionInventory, publicIPs=args.public)
+
+# query regions concurrently
+pool = ThreadPool(len(regionArgs))
+regionInventories = pool.map(invf, regionArgs)
+pool.close()
+pool.join()
+
+# combine regions
+inventory = reduce(mergeDictOfLists, regionInventories, {})
+inventory['_meta'] = {'hostvars': HOSTVARS}
+
+print(json.dumps(inventory))