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 and the IPs of the Instances within them, suitable
   5 for use as an Ansible inventory.
   6 '''
   7
   8 import argparse
   9 import boto3
  10 import json
  11 import sys
  12 from multiprocessing.dummy import Pool as ThreadPool
  13 from functools import partial
  14
  15
  16 DEFAULT_REGIONS = ['us-east-1', 'us-west-2']
  17 HOSTVARS = {}
  18
  19
  20 def allASGInstances(asgc):
  21     'Return a tuple of a dict of each ASG name listing the instance IDs within it, and a list of all instance IDs.'
  22     asgs = {}
  23     instanceIds = []
  24     args = {}
  25     while True:
  26         response = asgc.describe_auto_scaling_groups(**args)
  27         for asg in response['AutoScalingGroups']:
  28             asgs[asg['AutoScalingGroupName']] = [i['InstanceId'] for i in asg['Instances']]
  29             instanceIds += asgs[asg['AutoScalingGroupName']]
  30         if 'NextToken' not in response:
  31             break
  32         args = {'NextToken': response['NextToken']}
  33     return (asgs, instanceIds)
  34
  35
  36 def allInstanceIPs(ec2c, InstanceIds=None, publicIPs=False):
  37     'Return a dict of each Instance ID with its addresses.'
  38     instances = {}
  39     args = {}
  40     IPType = 'PublicIpAddress' if publicIPs else 'PrivateIpAddress'
  41     if InstanceIds is not None:
  42         args['InstanceIds'] = InstanceIds
  43     while True:
  44         response = ec2c.describe_instances(**args)
  45         for resv in response['Reservations']:
  46             for inst in resv['Instances']:
  47                 if IPType in inst:
  48                     instances[inst['InstanceId']] = inst[IPType]
  49         if 'NextToken' not in response:
  50             break
  51         args = {'NextToken': response['NextToken']}
  52     return instances
  53
  54
  55 def regionInventory(sessionArgs, publicIPs=False):
  56     'Return dict results for one region.'
  57     session = boto3.session.Session(**sessionArgs)
  58     asgc = session.client('autoscaling')
  59     ec2c = session.client('ec2')
  60
  61     # get dict of ASG names and associated Instance Ids, plus list of all Instance Ids referenced by ASGs
  62     (ASGs, AllInstanceIds) = allASGInstances(asgc)
  63
  64     # get list of instance IPs for all instance Ids used by ASGs
  65     AllInstanceIPs = allInstanceIPs(ec2c, InstanceIds=AllInstanceIds, publicIPs=publicIPs)
  66
  67     # a group for every Instance Id
  68     inventory = {iid:[AllInstanceIPs[iid]] for iid in AllInstanceIPs}
  69
  70     # add ASG dict, replacing ASG Instance Id with instance IP
  71     inventory.update({asg:[AllInstanceIPs[iid] for iid in ASGs[asg]] for asg in ASGs})
  72
  73     return inventory
  74
  75
  76 def mergeDictOfLists(a, b):
  77     'There is likely a better way of doing this, but right now I have a headache.'
  78     for key in b:
  79         if key in a:
  80             a[key] += b[key]
  81         else:
  82             a[key] = b[key]
  83     return a
  84
  85
  86 parser = argparse.ArgumentParser(description='dynamic Ansible inventory from AWS Autoscaling Groups')
  87 parser.add_argument('--public', action='store_true', help='inventory public IPs (default: private IPs)')
  88 parser.add_argument('--profile', metavar='PROFILE', dest='profile_name', help='AWS Profile (default: current IAM Role)')
  89 parser.add_argument('--regions', metavar='REGION', nargs='+', default=DEFAULT_REGIONS, help='AWS Regions (default: %(default)s)')
  90 parser.add_argument('--list', action='store_true')
  91 parser.add_argument('--host', nargs=1)
  92 args = parser.parse_args()
  93
  94 if args.host:
  95     print(json.dumps(HOSTVARS))
  96     sys.exit()
  97
  98 # create sessionArgs for each region
  99 regionArgs = [{'region_name': region} for region in args.regions]
 100 if args.profile_name:
 101     for arg in regionArgs:
 102         arg.update({'profile_name': args.profile_name})
 103
 104 # pin the non-variant option
 105 invf = partial(regionInventory, publicIPs=args.public)
 106
 107 # query regions concurrently
 108 pool = ThreadPool(len(regionArgs))
 109 regionInventories = pool.map(invf, regionArgs)
 110 pool.close()
 111 pool.join()
 112
 113 # combine regions
 114 inventory = reduce(mergeDictOfLists, regionInventories, {})
 115 inventory['_meta'] = {'hostvars': HOSTVARS}
 116
 117 print(json.dumps(inventory))