#!/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
import os
from multiprocessing.dummy import Pool as ThreadPool
from functools import partial


#DEFAULT_REGIONS = ['us-east-1', 'us-west-2']
DEFAULT_REGIONS = ['us-east-2']
HOSTVARS = {}

try:
    PUBLIC = len(os.environ['INVENTORY_PUBLIC']) > 0
except:
    PUBLIC = False

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] if iid in AllInstanceIPs] 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' if not PUBLIC else 'store_false', 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))