1#!/usr/bin/env python
2
3"""
4DigitalOcean external inventory script
5======================================
6
7Generates Ansible inventory of DigitalOcean Droplets.
8
9In addition to the --list and --host options used by Ansible, there are options
10for generating JSON of other DigitalOcean data.  This is useful when creating
11droplets.  For example, --regions will return all the DigitalOcean Regions.
12This information can also be easily found in the cache file, whose default
13location is /tmp/ansible-digital_ocean.cache).
14
15The --pretty (-p) option pretty-prints the output for better human readability.
16
17----
18Although the cache stores all the information received from DigitalOcean,
19the cache is not used for current droplet information (in --list, --host,
20--all, and --droplets).  This is so that accurate droplet information is always
21found.  You can force this script to use the cache with --force-cache.
22
23----
24Configuration is read from `digital_ocean.ini`, then from environment variables,
25and then from command-line arguments.
26
27Most notably, the DigitalOcean API Token must be specified. It can be specified
28in the INI file or with the following environment variables:
29    export DO_API_TOKEN='abc123' or
30    export DO_API_KEY='abc123'
31
32Alternatively, it can be passed on the command-line with --api-token.
33
34If you specify DigitalOcean credentials in the INI file, a handy way to
35get them into your environment (e.g., to use the digital_ocean module)
36is to use the output of the --env option with export:
37    export $(digital_ocean.py --env)
38
39----
40The following groups are generated from --list:
41 - ID    (droplet ID)
42 - NAME  (droplet NAME)
43 - digital_ocean
44 - image_ID
45 - image_NAME
46 - distro_NAME  (distribution NAME from image)
47 - region_NAME
48 - size_NAME
49 - status_STATUS
50
51For each host, the following variables are registered:
52 - do_backup_ids
53 - do_created_at
54 - do_disk
55 - do_features - list
56 - do_id
57 - do_image - object
58 - do_ip_address
59 - do_private_ip_address
60 - do_kernel - object
61 - do_locked
62 - do_memory
63 - do_name
64 - do_networks - object
65 - do_next_backup_window
66 - do_region - object
67 - do_size - object
68 - do_size_slug
69 - do_snapshot_ids - list
70 - do_status
71 - do_tags
72 - do_vcpus
73 - do_volume_ids
74
75-----
76```
77usage: digital_ocean.py [-h] [--list] [--host HOST] [--all] [--droplets]
78                        [--regions] [--images] [--sizes] [--ssh-keys]
79                        [--domains] [--tags] [--pretty]
80                        [--cache-path CACHE_PATH]
81                        [--cache-max_age CACHE_MAX_AGE] [--force-cache]
82                        [--refresh-cache] [--env] [--api-token API_TOKEN]
83
84Produce an Ansible Inventory file based on DigitalOcean credentials
85
86optional arguments:
87  -h, --help            show this help message and exit
88  --list                List all active Droplets as Ansible inventory
89                        (default: True)
90  --host HOST           Get all Ansible inventory variables about a specific
91                        Droplet
92  --all                 List all DigitalOcean information as JSON
93  --droplets, -d        List Droplets as JSON
94  --regions             List Regions as JSON
95  --images              List Images as JSON
96  --sizes               List Sizes as JSON
97  --ssh-keys            List SSH keys as JSON
98  --domains             List Domains as JSON
99  --tags                List Tags as JSON
100  --pretty, -p          Pretty-print results
101  --cache-path CACHE_PATH
102                        Path to the cache files (default: .)
103  --cache-max_age CACHE_MAX_AGE
104                        Maximum age of the cached items (default: 0)
105  --force-cache         Only use data from the cache
106  --refresh-cache, -r   Force refresh of cache by making API requests to
107                        DigitalOcean (default: False - use cache files)
108  --env, -e             Display DO_API_TOKEN
109  --api-token API_TOKEN, -a API_TOKEN
110                        DigitalOcean API Token
111```
112
113"""
114
115# (c) 2013, Evan Wies <evan@neomantra.net>
116# (c) 2017, Ansible Project
117# (c) 2017, Abhijeet Kasurde <akasurde@redhat.com>
118#
119# Inspired by the EC2 inventory plugin:
120# https://github.com/ansible/ansible/blob/devel/contrib/inventory/ec2.py
121#
122# This file is part of Ansible,
123#
124# Ansible is free software: you can redistribute it and/or modify
125# it under the terms of the GNU General Public License as published by
126# the Free Software Foundation, either version 3 of the License, or
127# (at your option) any later version.
128#
129# Ansible is distributed in the hope that it will be useful,
130# but WITHOUT ANY WARRANTY; without even the implied warranty of
131# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
132# GNU General Public License for more details.
133#
134# You should have received a copy of the GNU General Public License
135# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
136
137######################################################################
138
139import argparse
140import ast
141import os
142import re
143import requests
144import sys
145from time import time
146
147try:
148    import ConfigParser
149except ImportError:
150    import configparser as ConfigParser
151
152import json
153
154
155class DoManager:
156    def __init__(self, api_token):
157        self.api_token = api_token
158        self.api_endpoint = 'https://api.digitalocean.com/v2'
159        self.headers = {'Authorization': 'Bearer {0}'.format(self.api_token),
160                        'Content-type': 'application/json'}
161        self.timeout = 60
162
163    def _url_builder(self, path):
164        if path[0] == '/':
165            path = path[1:]
166        return '%s/%s' % (self.api_endpoint, path)
167
168    def send(self, url, method='GET', data=None):
169        url = self._url_builder(url)
170        data = json.dumps(data)
171        try:
172            if method == 'GET':
173                resp_data = {}
174                incomplete = True
175                while incomplete:
176                    resp = requests.get(url, data=data, headers=self.headers, timeout=self.timeout)
177                    json_resp = resp.json()
178
179                    for key, value in json_resp.items():
180                        if isinstance(value, list) and key in resp_data:
181                            resp_data[key] += value
182                        else:
183                            resp_data[key] = value
184
185                    try:
186                        url = json_resp['links']['pages']['next']
187                    except KeyError:
188                        incomplete = False
189
190        except ValueError as e:
191            sys.exit("Unable to parse result from %s: %s" % (url, e))
192        return resp_data
193
194    def all_active_droplets(self):
195        resp = self.send('droplets/')
196        return resp['droplets']
197
198    def all_regions(self):
199        resp = self.send('regions/')
200        return resp['regions']
201
202    def all_images(self, filter_name='global'):
203        params = {'filter': filter_name}
204        resp = self.send('images/', data=params)
205        return resp['images']
206
207    def sizes(self):
208        resp = self.send('sizes/')
209        return resp['sizes']
210
211    def all_ssh_keys(self):
212        resp = self.send('account/keys')
213        return resp['ssh_keys']
214
215    def all_domains(self):
216        resp = self.send('domains/')
217        return resp['domains']
218
219    def show_droplet(self, droplet_id):
220        resp = self.send('droplets/%s' % droplet_id)
221        return resp['droplet']
222
223    def all_tags(self):
224        resp = self.send('tags')
225        return resp['tags']
226
227
228class DigitalOceanInventory(object):
229
230    ###########################################################################
231    # Main execution path
232    ###########################################################################
233
234    def __init__(self):
235        """Main execution path """
236
237        # DigitalOceanInventory data
238        self.data = {}  # All DigitalOcean data
239        self.inventory = {}  # Ansible Inventory
240
241        # Define defaults
242        self.cache_path = '.'
243        self.cache_max_age = 0
244        self.use_private_network = False
245        self.group_variables = {}
246
247        # Read settings, environment variables, and CLI arguments
248        self.read_settings()
249        self.read_environment()
250        self.read_cli_args()
251
252        # Verify credentials were set
253        if not hasattr(self, 'api_token'):
254            msg = 'Could not find values for DigitalOcean api_token. They must be specified via either ini file, ' \
255                  'command line argument (--api-token), or environment variables (DO_API_TOKEN)\n'
256            sys.stderr.write(msg)
257            sys.exit(-1)
258
259        # env command, show DigitalOcean credentials
260        if self.args.env:
261            print("DO_API_TOKEN=%s" % self.api_token)
262            sys.exit(0)
263
264        # Manage cache
265        self.cache_filename = self.cache_path + "/ansible-digital_ocean.cache"
266        self.cache_refreshed = False
267
268        if self.is_cache_valid():
269            self.load_from_cache()
270            if len(self.data) == 0:
271                if self.args.force_cache:
272                    sys.stderr.write('Cache is empty and --force-cache was specified\n')
273                    sys.exit(-1)
274
275        self.manager = DoManager(self.api_token)
276
277        # Pick the json_data to print based on the CLI command
278        if self.args.droplets:
279            self.load_from_digital_ocean('droplets')
280            json_data = {'droplets': self.data['droplets']}
281        elif self.args.regions:
282            self.load_from_digital_ocean('regions')
283            json_data = {'regions': self.data['regions']}
284        elif self.args.images:
285            self.load_from_digital_ocean('images')
286            json_data = {'images': self.data['images']}
287        elif self.args.sizes:
288            self.load_from_digital_ocean('sizes')
289            json_data = {'sizes': self.data['sizes']}
290        elif self.args.ssh_keys:
291            self.load_from_digital_ocean('ssh_keys')
292            json_data = {'ssh_keys': self.data['ssh_keys']}
293        elif self.args.domains:
294            self.load_from_digital_ocean('domains')
295            json_data = {'domains': self.data['domains']}
296        elif self.args.tags:
297            self.load_from_digital_ocean('tags')
298            json_data = {'tags': self.data['tags']}
299        elif self.args.all:
300            self.load_from_digital_ocean()
301            json_data = self.data
302        elif self.args.host:
303            json_data = self.load_droplet_variables_for_host()
304        else:    # '--list' this is last to make it default
305            self.load_from_digital_ocean('droplets')
306            self.build_inventory()
307            json_data = self.inventory
308
309        if self.cache_refreshed:
310            self.write_to_cache()
311
312        if self.args.pretty:
313            print(json.dumps(json_data, indent=2))
314        else:
315            print(json.dumps(json_data))
316
317    ###########################################################################
318    # Script configuration
319    ###########################################################################
320
321    def read_settings(self):
322        """ Reads the settings from the digital_ocean.ini file """
323        config = ConfigParser.ConfigParser()
324        config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'digital_ocean.ini')
325        config.read(config_path)
326
327        # Credentials
328        if config.has_option('digital_ocean', 'api_token'):
329            self.api_token = config.get('digital_ocean', 'api_token')
330
331        # Cache related
332        if config.has_option('digital_ocean', 'cache_path'):
333            self.cache_path = config.get('digital_ocean', 'cache_path')
334        if config.has_option('digital_ocean', 'cache_max_age'):
335            self.cache_max_age = config.getint('digital_ocean', 'cache_max_age')
336
337        # Private IP Address
338        if config.has_option('digital_ocean', 'use_private_network'):
339            self.use_private_network = config.getboolean('digital_ocean', 'use_private_network')
340
341        # Group variables
342        if config.has_option('digital_ocean', 'group_variables'):
343            self.group_variables = ast.literal_eval(config.get('digital_ocean', 'group_variables'))
344
345    def read_environment(self):
346        """ Reads the settings from environment variables """
347        # Setup credentials
348        if os.getenv("DO_API_TOKEN"):
349            self.api_token = os.getenv("DO_API_TOKEN")
350        if os.getenv("DO_API_KEY"):
351            self.api_token = os.getenv("DO_API_KEY")
352
353    def read_cli_args(self):
354        """ Command line argument processing """
355        parser = argparse.ArgumentParser(description='Produce an Ansible Inventory file based on DigitalOcean credentials')
356
357        parser.add_argument('--list', action='store_true', help='List all active Droplets as Ansible inventory (default: True)')
358        parser.add_argument('--host', action='store', help='Get all Ansible inventory variables about a specific Droplet')
359
360        parser.add_argument('--all', action='store_true', help='List all DigitalOcean information as JSON')
361        parser.add_argument('--droplets', '-d', action='store_true', help='List Droplets as JSON')
362        parser.add_argument('--regions', action='store_true', help='List Regions as JSON')
363        parser.add_argument('--images', action='store_true', help='List Images as JSON')
364        parser.add_argument('--sizes', action='store_true', help='List Sizes as JSON')
365        parser.add_argument('--ssh-keys', action='store_true', help='List SSH keys as JSON')
366        parser.add_argument('--domains', action='store_true', help='List Domains as JSON')
367        parser.add_argument('--tags', action='store_true', help='List Tags as JSON')
368
369        parser.add_argument('--pretty', '-p', action='store_true', help='Pretty-print results')
370
371        parser.add_argument('--cache-path', action='store', help='Path to the cache files (default: .)')
372        parser.add_argument('--cache-max_age', action='store', help='Maximum age of the cached items (default: 0)')
373        parser.add_argument('--force-cache', action='store_true', default=False, help='Only use data from the cache')
374        parser.add_argument('--refresh-cache', '-r', action='store_true', default=False,
375                            help='Force refresh of cache by making API requests to DigitalOcean (default: False - use cache files)')
376
377        parser.add_argument('--env', '-e', action='store_true', help='Display DO_API_TOKEN')
378        parser.add_argument('--api-token', '-a', action='store', help='DigitalOcean API Token')
379
380        self.args = parser.parse_args()
381
382        if self.args.api_token:
383            self.api_token = self.args.api_token
384
385        # Make --list default if none of the other commands are specified
386        if (not self.args.droplets and not self.args.regions and
387                not self.args.images and not self.args.sizes and
388                not self.args.ssh_keys and not self.args.domains and
389                not self.args.tags and
390                not self.args.all and not self.args.host):
391            self.args.list = True
392
393    ###########################################################################
394    # Data Management
395    ###########################################################################
396
397    def load_from_digital_ocean(self, resource=None):
398        """Get JSON from DigitalOcean API """
399        if self.args.force_cache and os.path.isfile(self.cache_filename):
400            return
401        # We always get fresh droplets
402        if self.is_cache_valid() and not (resource == 'droplets' or resource is None):
403            return
404        if self.args.refresh_cache:
405            resource = None
406
407        if resource == 'droplets' or resource is None:
408            self.data['droplets'] = self.manager.all_active_droplets()
409            self.cache_refreshed = True
410        if resource == 'regions' or resource is None:
411            self.data['regions'] = self.manager.all_regions()
412            self.cache_refreshed = True
413        if resource == 'images' or resource is None:
414            self.data['images'] = self.manager.all_images()
415            self.cache_refreshed = True
416        if resource == 'sizes' or resource is None:
417            self.data['sizes'] = self.manager.sizes()
418            self.cache_refreshed = True
419        if resource == 'ssh_keys' or resource is None:
420            self.data['ssh_keys'] = self.manager.all_ssh_keys()
421            self.cache_refreshed = True
422        if resource == 'domains' or resource is None:
423            self.data['domains'] = self.manager.all_domains()
424            self.cache_refreshed = True
425        if resource == 'tags' or resource is None:
426            self.data['tags'] = self.manager.all_tags()
427            self.cache_refreshed = True
428
429    def add_inventory_group(self, key):
430        """ Method to create group dict """
431        host_dict = {'hosts': [], 'vars': {}}
432        self.inventory[key] = host_dict
433        return
434
435    def add_host(self, group, host):
436        """ Helper method to reduce host duplication """
437        if group not in self.inventory:
438            self.add_inventory_group(group)
439
440        if host not in self.inventory[group]['hosts']:
441            self.inventory[group]['hosts'].append(host)
442        return
443
444    def build_inventory(self):
445        """ Build Ansible inventory of droplets """
446        self.inventory = {
447            'all': {
448                'hosts': [],
449                'vars': self.group_variables
450            },
451            '_meta': {'hostvars': {}}
452        }
453
454        # add all droplets by id and name
455        for droplet in self.data['droplets']:
456            for net in droplet['networks']['v4']:
457                if net['type'] == 'public':
458                    dest = net['ip_address']
459                else:
460                    continue
461
462            self.inventory['all']['hosts'].append(dest)
463
464            self.add_host(droplet['id'], dest)
465
466            self.add_host(droplet['name'], dest)
467
468            # groups that are always present
469            for group in ('digital_ocean',
470                          'region_' + droplet['region']['slug'],
471                          'image_' + str(droplet['image']['id']),
472                          'size_' + droplet['size']['slug'],
473                          'distro_' + DigitalOceanInventory.to_safe(droplet['image']['distribution']),
474                          'status_' + droplet['status']):
475                self.add_host(group, dest)
476
477            # groups that are not always present
478            for group in (droplet['image']['slug'],
479                          droplet['image']['name']):
480                if group:
481                    image = 'image_' + DigitalOceanInventory.to_safe(group)
482                    self.add_host(image, dest)
483
484            if droplet['tags']:
485                for tag in droplet['tags']:
486                    self.add_host(tag, dest)
487
488            # hostvars
489            info = self.do_namespace(droplet)
490            self.inventory['_meta']['hostvars'][dest] = info
491
492    def load_droplet_variables_for_host(self):
493        """ Generate a JSON response to a --host call """
494        host = int(self.args.host)
495        droplet = self.manager.show_droplet(host)
496        info = self.do_namespace(droplet)
497        return {'droplet': info}
498
499    ###########################################################################
500    # Cache Management
501    ###########################################################################
502
503    def is_cache_valid(self):
504        """ Determines if the cache files have expired, or if it is still valid """
505        if os.path.isfile(self.cache_filename):
506            mod_time = os.path.getmtime(self.cache_filename)
507            current_time = time()
508            if (mod_time + self.cache_max_age) > current_time:
509                return True
510        return False
511
512    def load_from_cache(self):
513        """ Reads the data from the cache file and assigns it to member variables as Python Objects """
514        try:
515            with open(self.cache_filename, 'r') as cache:
516                json_data = cache.read()
517            data = json.loads(json_data)
518        except IOError:
519            data = {'data': {}, 'inventory': {}}
520
521        self.data = data['data']
522        self.inventory = data['inventory']
523
524    def write_to_cache(self):
525        """ Writes data in JSON format to a file """
526        data = {'data': self.data, 'inventory': self.inventory}
527        json_data = json.dumps(data, indent=2)
528
529        with open(self.cache_filename, 'w') as cache:
530            cache.write(json_data)
531
532    ###########################################################################
533    # Utilities
534    ###########################################################################
535    @staticmethod
536    def to_safe(word):
537        """ Converts 'bad' characters in a string to underscores so they can be used as Ansible groups """
538        return re.sub(r"[^A-Za-z0-9\-.]", "_", word)
539
540    @staticmethod
541    def do_namespace(data):
542        """ Returns a copy of the dictionary with all the keys put in a 'do_' namespace """
543        info = {}
544        for k, v in data.items():
545            info['do_' + k] = v
546        return info
547
548
549###########################################################################
550# Run the script
551DigitalOceanInventory()
552