git.squeep.com Git - awsible/blob - inventory/asg-inventory.py

   1 #!/usr/bin/env python
   2 '''\
   3 Generate a JSON object containing the names of all the AWS Autoscaling
   4 Groups in an account, the IDs of the Instances within them, as well
   5 as all the Tags on those Instances, each containing a list of the IPs
   6 of the Instances, suitable for use as an Ansible inventory.
   7
   8 This output is similar to the default ec2.py inventory script, but is
   9 far more limited in the scope of information is has to retreive, parse,
  10 and render.
  11
  12 It does not currently deal with hostvars at all.
  13 '''
  14
  15 import argparse
  16 import boto3
  17 import json
  18 import sys
  19 import os
  20 from multiprocessing.dummy import Pool as ThreadPool
  21 from functools import partial
  22
  23
  24 #DEFAULT_REGIONS = ['us-east-1', 'us-west-2']
  25 DEFAULT_REGIONS = ['us-east-2']
  26 HOSTVARS = {}
  27
  28 try:
  29     PUBLIC = len(os.environ['INVENTORY_PUBLIC']) > 0
  30 except:
  31     PUBLIC = False
  32
  33 def allASGInstances(asgc):
  34     'Return a tuple of a dict of each ASG name listing the instance IDs within it, and a list of all instance IDs.'
  35     asgs = {}
  36     instanceIds = []
  37     args = {}
  38     while True:
  39         response = asgc.describe_auto_scaling_groups(**args)
  40         for asg in response['AutoScalingGroups']:
  41             asgs[asg['AutoScalingGroupName']] = [i['InstanceId'] for i in asg['Instances']]
  42             instanceIds += asgs[asg['AutoScalingGroupName']]
  43         if 'NextToken' not in response:
  44             break
  45         args = {'NextToken': response['NextToken']}
  46     return (asgs, instanceIds)
  47
  48
  49 def allInstanceIPs(ec2c, InstanceIds=None, publicIPs=False):
  50     'Return a dict of each Instance ID with its addresses.'
  51     instances = {}
  52     args = {}
  53     IPType = 'PublicIpAddress' if publicIPs else 'PrivateIpAddress'
  54     if InstanceIds is not None:
  55         args['InstanceIds'] = InstanceIds
  56     while True:
  57         response = ec2c.describe_instances(**args)
  58         for resv in response['Reservations']:
  59             for inst in resv['Instances']:
  60                 if IPType in inst:
  61                     instances[inst['InstanceId']] = inst[IPType]
  62         if 'NextToken' not in response:
  63             break
  64         args = {'NextToken': response['NextToken']}
  65     return instances
  66
  67
  68 def tagsOfInstances(ec2c, InstanceIds=None):
  69     tags = {}
  70     args = {'Filters': [{'Name': 'resource-type', 'Values': ["instance"]}]}
  71     if InstanceIds:
  72         args['Filters'] += [{'Name': 'resource-id', 'Values': InstanceIds}]
  73     while True:
  74         response = ec2c.describe_tags(**args)
  75         for tag in response['Tags']:
  76             tagname = "tag_{}_{}".format(tag['Key'], tag['Value'])
  77             tagvalue = tag['ResourceId']
  78             tags.setdefault(tagname, []).append(tagvalue)
  79         if 'NextToken' not in response:
  80             break
  81         args['NextToken'] = response['NextToken']
  82     return tags
  83
  84
  85 def regionInventory(sessionArgs, publicIPs=False):
  86     'Return dict results for one region.'
  87     session = boto3.session.Session(**sessionArgs)
  88     asgc = session.client('autoscaling')
  89     ec2c = session.client('ec2')
  90
  91     # get dict of ASG names and associated Instance Ids, plus list of all Instance Ids referenced by ASGs
  92     (ASGs, AllInstanceIds) = allASGInstances(asgc)
  93
  94     # get list of instance IPs for all instance Ids used by ASGs
  95     AllInstanceIPs = allInstanceIPs(ec2c, InstanceIds=AllInstanceIds, publicIPs=publicIPs)
  96
  97     # a group for every Instance Id
  98     inventory = {iid:[AllInstanceIPs[iid]] for iid in AllInstanceIPs}
  99
 100     # add ASG dict, replacing ASG Instance Id with instance IP
 101     inventory.update({asg:[AllInstanceIPs[iid] for iid in ASGs[asg] if iid in AllInstanceIPs] for asg in ASGs})
 102
 103     # group up instance tags as well
 104     tags = tagsOfInstances(ec2c, AllInstanceIds)
 105     inventory.update({tag:[AllInstanceIPs[iid] for iid in tags[tag] if iid in AllInstanceIPs] for tag in tags})
 106
 107     return inventory
 108
 109
 110 def mergeDictOfLists(a, b):
 111     'There is likely a better way of doing this, but right now I have a headache.'
 112     for key in b:
 113         if key in a:
 114             a[key] += b[key]
 115         else:
 116             a[key] = b[key]
 117     return a
 118
 119
 120 parser = argparse.ArgumentParser(description='dynamic Ansible inventory from AWS Autoscaling Groups')
 121 parser.add_argument('--public', action='store_true' if not PUBLIC else 'store_false', help='inventory public IPs (default: private IPs)')
 122 parser.add_argument('--profile', metavar='PROFILE', dest='profile_name', help='AWS Profile (default: current IAM Role)')
 123 parser.add_argument('--regions', metavar='REGION', nargs='+', default=DEFAULT_REGIONS, help='AWS Regions (default: %(default)s)')
 124 parser.add_argument('--list', action='store_true')
 125 parser.add_argument('--host', nargs=1)
 126 args = parser.parse_args()
 127
 128 if args.host:
 129     print(json.dumps(HOSTVARS))
 130     sys.exit()
 131
 132 # create sessionArgs for each region
 133 regionArgs = [{'region_name': region} for region in args.regions]
 134 if args.profile_name:
 135     for arg in regionArgs:
 136         arg.update({'profile_name': args.profile_name})
 137
 138 # pin the non-variant option
 139 invf = partial(regionInventory, publicIPs=args.public)
 140
 141 # query regions concurrently
 142 pool = ThreadPool(len(regionArgs))
 143 regionInventories = pool.map(invf, regionArgs)
 144 pool.close()
 145 pool.join()
 146
 147 # combine regions
 148 inventory = reduce(mergeDictOfLists, regionInventories, {})
 149 inventory['_meta'] = {'hostvars': HOSTVARS}
 150
 151 print(json.dumps(inventory))