/plugins/inventory/collins.py

  1#!/usr/bin/env python
  2
  3"""
  4Collins external inventory script
  5=================================
  6
  7Ansible has a feature where instead of reading from /etc/ansible/hosts
  8as a text file, it can query external programs to obtain the list
  9of hosts, groups the hosts are in, and even variables to assign to each host.
 10
 11Collins is a hardware asset management system originally developed by
 12Tumblr for tracking new hardware as it built out its own datacenters.  It
 13exposes a rich API for manipulating and querying one's hardware inventory,
 14which makes it an ideal 'single point of truth' for driving systems
 15automation like Ansible.  Extensive documentation on Collins, including a quickstart,
 16API docs, and a full reference manual, can be found here:
 17
 18http://tumblr.github.io/collins
 19
 20This script adds support to Ansible for obtaining a dynamic inventory of
 21assets in your infrastructure, grouping them in Ansible by their useful attributes,
 22and binding all facts provided by Collins to each host so that they can be used to
 23drive automation.  Some parts of this script were cribbed shamelessly from mdehaan's
 24Cobbler inventory script.
 25
 26To use it, copy it to your repo and pass -i <collins script> to the ansible or
 27ansible-playbook command; if you'd like to use it by default, simply copy collins.ini
 28to /etc/ansible and this script to /etc/ansible/hosts.
 29
 30Alongside the options set in collins.ini, there are several environment variables
 31that will be used instead of the configured values if they are set:
 32
 33 - COLLINS_USERNAME - specifies a username to use for Collins authentication
 34 - COLLINS_PASSWORD - specifies a password to use for Collins authentication
 35 - COLLINS_ASSET_TYPE - specifies a Collins asset type to use during querying;
 36   this can be used to run Ansible automation against different asset classes than
 37   server nodes, such as network switches and PDUs
 38 - COLLINS_CONFIG - specifies an alternative location for collins.ini, defaults to
 39   <location of collins.py>/collins.ini
 40
 41If errors are encountered during operation, this script will return an exit code of
 42255; otherwise, it will return an exit code of 0.
 43
 44Tested against Ansible 1.6.6 and Collins 1.2.4.
 45"""
 46
 47# (c) 2014, Steve Salevan <steve.salevan@gmail.com>
 48#
 49# This file is part of Ansible.
 50#
 51# Ansible is free software: you can redistribute it and/or modify
 52# it under the terms of the GNU General Public License as published by
 53# the Free Software Foundation, either version 3 of the License, or
 54# (at your option) any later version.
 55#
 56# Ansible is distributed in the hope that it will be useful,
 57# but WITHOUT ANY WARRANTY; without even the implied warranty of
 58# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 59# GNU General Public License for more details.
 60#
 61# You should have received a copy of the GNU General Public License
 62# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
 63
 64######################################################################
 65
 66
 67import argparse
 68import base64
 69import ConfigParser
 70import logging
 71import os
 72import re
 73import sys
 74from time import time
 75import traceback
 76import urllib
 77import urllib2
 78
 79try:
 80    import json
 81except ImportError:
 82    import simplejson as json
 83
 84
 85class CollinsDefaults(object):
 86    ASSETS_API_ENDPOINT = '%s/api/assets'
 87    SPECIAL_ATTRIBUTES = set([
 88        'CREATED',
 89        'DELETED',
 90        'UPDATED',
 91        'STATE',
 92    ])
 93    LOG_FORMAT = '%(asctime)-15s %(message)s'
 94
 95
 96class Error(Exception):
 97    pass
 98
 99
100class MaxRetriesError(Error):
101    pass
102
103
104class CollinsInventory(object):
105
106    def __init__(self):
107        """ Constructs CollinsInventory object and reads all configuration. """
108
109        self.inventory = dict()  # A list of groups and the hosts in that group
110        self.cache = dict()  # Details about hosts in the inventory
111
112        # Read settings and parse CLI arguments
113        self.read_settings()
114        self.parse_cli_args()
115
116        logging.basicConfig(format=CollinsDefaults.LOG_FORMAT,
117            filename=self.log_location)
118        self.log = logging.getLogger('CollinsInventory')
119
120    def _asset_get_attribute(self, asset, attrib):
121        """ Returns a user-defined attribute from an asset if it exists; otherwise,
122            returns None. """
123
124        if 'ATTRIBS' in asset:
125            for attrib_block in asset['ATTRIBS'].keys():
126                if attrib in asset['ATTRIBS'][attrib_block]:
127                    return asset['ATTRIBS'][attrib_block][attrib]
128        return None
129
130    def _asset_has_attribute(self, asset, attrib):
131        """ Returns whether a user-defined attribute is present on an asset. """
132
133        if 'ATTRIBS' in asset:
134            for attrib_block in asset['ATTRIBS'].keys():
135                if attrib in asset['ATTRIBS'][attrib_block]:
136                    return True
137        return False
138
139    def run(self):
140        """ Main execution path """
141
142        # Updates cache if cache is not present or has expired.
143        successful = True
144        if self.args.refresh_cache:
145            successful = self.update_cache()
146        elif not self.is_cache_valid():
147            successful = self.update_cache()
148        else:
149            successful = self.load_inventory_from_cache()
150            successful &= self.load_cache_from_cache()
151
152        data_to_print = ""
153
154        # Data to print
155        if self.args.host:
156            data_to_print = self.get_host_info()
157
158        elif self.args.list:
159            # Display list of instances for inventory
160            data_to_print = self.json_format_dict(self.inventory, self.args.pretty)
161
162        else:  # default action with no options
163            data_to_print = self.json_format_dict(self.inventory, self.args.pretty)
164
165        print data_to_print
166        return successful
167
168    def find_assets(self, attributes = {}, operation = 'AND'):
169        """ Obtains Collins assets matching the provided attributes. """
170
171        # Formats asset search query to locate assets matching attributes, using
172        # the CQL search feature as described here:
173        # http://tumblr.github.io/collins/recipes.html
174        attributes_query = [ '='.join(attr_pair)
175            for attr_pair in attributes.iteritems() ]
176        query_parameters = {
177            'details': ['True'],
178            'operation': [operation],
179            'query': attributes_query,
180            'remoteLookup': [str(self.query_remote_dcs)],
181            'size': [self.results_per_query],
182            'type': [self.collins_asset_type],
183        }
184        assets = []
185        cur_page = 0
186        num_retries = 0
187        # Locates all assets matching the provided query, exhausting pagination.
188        while True:
189            if num_retries == self.collins_max_retries:
190                raise MaxRetriesError("Maximum of %s retries reached; giving up" % \
191                    self.collins_max_retries)
192            query_parameters['page'] = cur_page
193            query_url = "%s?%s" % (
194                (CollinsDefaults.ASSETS_API_ENDPOINT % self.collins_host),
195                urllib.urlencode(query_parameters, doseq=True)
196            )
197            request = urllib2.Request(query_url)
198            request.add_header('Authorization', self.basic_auth_header)
199            try:
200                response = urllib2.urlopen(request, timeout=self.collins_timeout_secs)
201                json_response = json.loads(response.read())
202                # Adds any assets found to the array of assets.
203                assets += json_response['data']['Data']
204                # If we've retrieved all of our assets, breaks out of the loop.
205                if len(json_response['data']['Data']) == 0:
206                    break
207                cur_page += 1
208                num_retries = 0
209            except:
210                self.log.error("Error while communicating with Collins, retrying:\n%s",
211                    traceback.format_exc())
212                num_retries += 1
213        return assets
214
215    def is_cache_valid(self):
216        """ Determines if the cache files have expired, or if it is still valid """
217
218        if os.path.isfile(self.cache_path_cache):
219            mod_time = os.path.getmtime(self.cache_path_cache)
220            current_time = time()
221            if (mod_time + self.cache_max_age) > current_time:
222                if os.path.isfile(self.cache_path_inventory):
223                    return True
224
225        return False
226
227    def read_settings(self):
228        """ Reads the settings from the collins.ini file """
229
230        config_loc = os.getenv('COLLINS_CONFIG',
231            os.path.dirname(os.path.realpath(__file__)) + '/collins.ini')
232
233        config = ConfigParser.SafeConfigParser()
234        config.read(os.path.dirname(os.path.realpath(__file__)) + '/collins.ini')
235
236        self.collins_host = config.get('collins', 'host')
237        self.collins_username = os.getenv('COLLINS_USERNAME',
238            config.get('collins', 'username'))
239        self.collins_password = os.getenv('COLLINS_PASSWORD',
240            config.get('collins', 'password'))
241        self.collins_asset_type = os.getenv('COLLINS_ASSET_TYPE',
242            config.get('collins', 'asset_type'))
243        self.collins_timeout_secs = config.getint('collins', 'timeout_secs')
244        self.collins_max_retries = config.getint('collins', 'max_retries')
245
246        self.results_per_query = config.getint('collins', 'results_per_query')
247        self.ip_address_index = config.getint('collins', 'ip_address_index')
248        self.query_remote_dcs = config.getboolean('collins', 'query_remote_dcs')
249        self.prefer_hostnames = config.getboolean('collins', 'prefer_hostnames')
250
251        cache_path = config.get('collins', 'cache_path')
252        self.cache_path_cache = cache_path + \
253            '/ansible-collins-%s.cache' % self.collins_asset_type
254        self.cache_path_inventory = cache_path + \
255            '/ansible-collins-%s.index' % self.collins_asset_type
256        self.cache_max_age = config.getint('collins', 'cache_max_age')
257
258        log_path = config.get('collins', 'log_path')
259        self.log_location = log_path + '/ansible-collins.log'
260        self.basic_auth_header = "Basic %s" % base64.encodestring(
261            '%s:%s' % (self.collins_username, self.collins_password))[:-1]
262
263    def parse_cli_args(self):
264        """ Command line argument processing """
265
266        parser = argparse.ArgumentParser(
267            description='Produces an Ansible Inventory file based on Collins')
268        parser.add_argument('--list',
269            action='store_true', default=True, help='List instances (default: True)')
270        parser.add_argument('--host',
271            action='store', help='Get all the variables about a specific instance')
272        parser.add_argument('--refresh-cache',
273            action='store_true', default=False,
274            help='Force refresh of cache by making API requests to Collins ' \
275                 '(default: False - use cache files)')
276        parser.add_argument('--pretty',
277            action='store_true', default=False, help='Pretty print all JSON output')
278        self.args = parser.parse_args()
279
280    def update_cache(self):
281        """ Make calls to Collins and saves the output in a cache """
282
283        self.cache = dict()
284        self.inventory = dict()
285
286        # Locates all server assets from Collins.
287        try:
288            server_assets = self.find_assets()
289        except:
290            self.log.error("Error while locating assets from Collins:\n%s",
291                traceback.format_exc())
292            return False
293
294        for asset in server_assets:
295            # Determines the index to retrieve the asset's IP address either by an
296            # attribute set on the Collins asset or the pre-configured value.
297            if self._asset_has_attribute(asset, 'ANSIBLE_IP_INDEX'):
298                ip_index = self._asset_get_attribute(asset, 'ANSIBLE_IP_INDEX')
299                try:
300                    ip_index = int(ip_index)
301                except:
302                    self.log.error(
303                        "ANSIBLE_IP_INDEX attribute on asset %s not an integer: %s", asset,
304                        ip_index)
305            else:
306                ip_index = self.ip_address_index
307
308            # Attempts to locate the asset's primary identifier (hostname or IP address),
309            # which will be used to index the asset throughout the Ansible inventory.
310            if self.prefer_hostnames and self._asset_has_attribute(asset, 'HOSTNAME'):
311                asset_identifier = self._asset_get_attribute(asset, 'HOSTNAME')
312            elif 'ADDRESSES' not in asset:
313                self.log.warning("No IP addresses found for asset '%s', skipping",
314                    asset)
315                continue
316            elif len(asset['ADDRESSES']) < ip_index + 1:
317                self.log.warning(
318                    "No IP address found at index %s for asset '%s', skipping",
319                    ip_index, asset)
320                continue
321            else:
322                asset_identifier = asset['ADDRESSES'][ip_index]['ADDRESS']
323
324            # Adds an asset index to the Ansible inventory based upon unpacking
325            # the name of the asset's current STATE from its dictionary.
326            if 'STATE' in asset['ASSET'] and asset['ASSET']['STATE']:
327                state_inventory_key = self.to_safe(
328                    'STATE-%s' % asset['ASSET']['STATE']['NAME'])
329                self.push(self.inventory, state_inventory_key, asset_identifier)
330
331            # Indexes asset by all user-defined Collins attributes.
332            if 'ATTRIBS' in asset:
333                for attrib_block in asset['ATTRIBS'].keys():
334                    for attrib in asset['ATTRIBS'][attrib_block].keys():
335                        attrib_key = self.to_safe(
336                            '%s-%s' % (attrib, asset['ATTRIBS'][attrib_block][attrib]))
337                        self.push(self.inventory, attrib_key, asset_identifier)
338
339            # Indexes asset by all built-in Collins attributes.
340            for attribute in asset['ASSET'].keys():
341                if attribute not in CollinsDefaults.SPECIAL_ATTRIBUTES:
342                    attribute_val = asset['ASSET'][attribute]
343                    if attribute_val is not None:
344                        attrib_key = self.to_safe('%s-%s' % (attribute, attribute_val))
345                        self.push(self.inventory, attrib_key, asset_identifier)
346
347            # Indexes asset by hardware product information.
348            if 'HARDWARE' in asset:
349                if 'PRODUCT' in asset['HARDWARE']['BASE']:
350                    product = asset['HARDWARE']['BASE']['PRODUCT']
351                    if product:
352                        product_key = self.to_safe(
353                            'HARDWARE-PRODUCT-%s' % asset['HARDWARE']['BASE']['PRODUCT'])
354                        self.push(self.inventory, product_key, asset_identifier)
355
356            # Indexing now complete, adds the host details to the asset cache.
357            self.cache[asset_identifier] = asset
358
359        try:
360            self.write_to_cache(self.cache, self.cache_path_cache)
361            self.write_to_cache(self.inventory, self.cache_path_inventory)
362        except:
363            self.log.error("Error while writing to cache:\n%s", traceback.format_exc())
364            return False
365        return True
366
367    def push(self, dictionary, key, value):
368        """ Adds a value to a list at a dictionary key, creating the list if it doesn't
369            exist. """
370
371        if key not in dictionary:
372            dictionary[key] = []
373        dictionary[key].append(value)
374
375    def get_host_info(self):
376        """ Get variables about a specific host. """
377
378        if not self.cache or len(self.cache) == 0:
379            # Need to load index from cache
380            self.load_cache_from_cache()
381
382        if not self.args.host in self.cache:
383            # try updating the cache
384            self.update_cache()
385
386            if not self.args.host in self.cache:
387                # host might not exist anymore
388                return self.json_format_dict({}, self.args.pretty)
389
390        return self.json_format_dict(self.cache[self.args.host], self.args.pretty)
391
392    def load_inventory_from_cache(self):
393        """ Reads the index from the cache file sets self.index """
394
395        try:
396            cache = open(self.cache_path_inventory, 'r')
397            json_inventory = cache.read()
398            self.inventory = json.loads(json_inventory)
399            return True
400        except:
401            self.log.error("Error while loading inventory:\n%s",
402                traceback.format_exc())
403            self.inventory = {}
404            return False
405
406    def load_cache_from_cache(self):
407        """ Reads the cache from the cache file sets self.cache """
408
409        try:
410            cache = open(self.cache_path_cache, 'r')
411            json_cache = cache.read()
412            self.cache = json.loads(json_cache)
413            return True
414        except:
415            self.log.error("Error while loading host cache:\n%s",
416                traceback.format_exc())
417            self.cache = {}
418            return False
419
420    def write_to_cache(self, data, filename):
421        """ Writes data in JSON format to a specified file. """
422
423        json_data = self.json_format_dict(data, self.args.pretty)
424        cache = open(filename, 'w')
425        cache.write(json_data)
426        cache.close()
427
428    def to_safe(self, word):
429        """ Converts 'bad' characters in a string to underscores so they
430            can be used as Ansible groups """
431
432        return re.sub("[^A-Za-z0-9\-]", "_", word)
433
434    def json_format_dict(self, data, pretty=False):
435        """ Converts a dict to a JSON object and dumps it as a formatted string """
436
437        if pretty:
438            return json.dumps(data, sort_keys=True, indent=2)
439        else:
440            return json.dumps(data)
441
442
443if __name__ in '__main__':
444    inventory = CollinsInventory()
445    if inventory.run():
446        sys.exit(0)
447    else:
448        sys.exit(-1)