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