AWS ElasticSearch service Ansible dynamic inventory

elasticsearch.py

#!/usr/bin/env python3

import argparse
import json
import re
import boto3
from botocore.config import Config

try:
    from collections.abc import Mapping
except ImportError:
    from collections import Mapping

REGIONS = ['us-east-1', 'us-west-2', 'eu-west-1']

class ElasticSearchInventory:
    def _empty_inventory(self):
        return {"_meta": {"hostvars": {}}}

    def __init__(self):
        ''' Main execution path '''
        self.inventory = self._empty_inventory()

        # Parse CLI arguments
        self.parse_cli_args()
        self.build_inventory()

        # Data to print
        if self.args.host:
            data_to_print = self.get_host_info()

        elif self.args.list:
            # Display list of instances for inventory
            data_to_print = self.inventory

        print(json.dumps(data_to_print, sort_keys=True, default=str))

    def parse_cli_args(self):
        ''' Command line argument processing '''

        parser = argparse.ArgumentParser(description='ElasticSearch Ansible Inventory')
        parser.add_argument('--list', action='store_true', default=True,
                            help='List instances (default: True)')
        parser.add_argument('--host', action='store',
                            help='Get all the variables about a specific instance')
        self.args = parser.parse_args()

    def build_inventory(self):
        for region in REGIONS:
            self.build_inventory_for_region(region)

    def build_inventory_for_region(self, region):
        client = boto3.client('es', config=Config(region_name=region))

        for dn in client.list_domain_names()['DomainNames']:
            domain_name = dn['DomainName']
            domain_info = client.describe_elasticsearch_domain(DomainName=domain_name)['DomainStatus']
            domain_tags = client.list_tags(ARN=domain_info['ARN'])['TagList']

            self.add_elasticsearch_domain(domain_info, domain_tags)

    def add_elasticsearch_domain(self, domain_info, domain_tags):
        domain_info = self.uncamelize_dict(domain_info)
        domain_info['tags'] = {tag['Key']: tag['Value'] for tag in domain_tags}
        domain_name = domain_info['domain_name']

        self.inventory['_meta']['hostvars'][domain_name] = domain_info
        
        # Add to general group
        self.add_to_group('es', domain_name)

        # Group by tags
        for key, value in domain_info['tags'].items():
            self.add_to_group('es_' + key + '_' + value, domain_name)

    def add_to_group(self, group_name, host_name):
        group_name = self.to_safe(group_name)
        self.inventory.setdefault(group_name, {'hosts': []})
        self.inventory[group_name]['hosts'].append(host_name)

    def get_host_info(self):
        return self.inventory['_meta']['hostvars'].get(self.args.host, {})

    def uncamelize_dict(self, dict):
        result = {}

        for key, value in dict.items():
            new_key = self.uncamelize(key)

            if isinstance(value, Mapping):
                result[new_key] = self.uncamelize_dict(value)
            else:
                result[new_key] = value

        return result

    def uncamelize(self, key):
        temp = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', key)
        return re.sub('([a-z0-9])([A-Z])', r'\1_\2', temp).lower()

    def to_safe(self, word):
        ''' Converts 'bad' characters in a string to underscores so they can be used as Ansible groups '''
        return re.sub('[^A-Za-z0-9\_]', '_', word)

if __name__ == '__main__':
    # Run the script
    ElasticSearchInventory()