/plugins/inventory/collins.py

https://github.com/ajanthanm/ansible · Python · 448 lines · 350 code · 31 blank · 67 comment · 38 complexity · 7d0c458a6ccf8bbcc22c2cbdda7a4195 MD5 · raw file

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