ansible/cloud/digital_ocean_droplet

#!/usr/bin/python
# -*- coding: utf-8 -*-

# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.
DOCUMENTATION = '''
---
module: digital_ocean_droplet
short_description: Create/delete a droplet in DigitalOcean
description:
     - Create/delete a droplet in DigitalOcean and optionally waits for it to be 'running'.
version_added: "1.6"
options:
  state:
    description:
     - Indicate desired state of the target.
    default: present
    choices: ['present', 'absent']
  client_id:
     description:
     - Digital Ocean manager id.
  api_key:
    description:
     - Digital Ocean api key.
  id:
    description:
     - Numeric, the droplet id you want to operate on.
  name:
    description:
     - String, this is the name of the droplet - must be formatted by hostname rules.
  unique_name:
    description:
     - Bool, require unique hostnames.  By default, digital ocean allows multiple hosts with the same name.  Setting this to "yes" allows only one host per name.  Useful for idempotence.
    default: "no"
    choices: [ "yes", "no" ]
  size_id:
    description:
     - Numeric, this is the id of the size you would like the droplet created at.
  image_id:
    description:
     - Numeric, this is the id of the image you would like the droplet created with.
  region_id:
    description:
     - "Numeric, this is the id of the region you would like your server"
  ssh_key_ids:
    description:
     - Optional, comma separated list of ssh_key_ids that you would like to be added to the server
  wait:
    description:
     - Wait for the droplet to be in state 'running' before returning.  If wait is "no" an ip_address may not be returned.
    default: "yes"
    choices: [ "yes", "no" ]
  wait_timeout:
    description:
     - How long before wait gives up, in seconds.
    default: 300

notes:
  - Two environment variables can be used, DO_CLIENT_ID and DO_API_KEY.
'''

EXAMPLES = '''
# Create a new Droplet
# Will return the droplet details including the droplet id (used for idempotence)

- digital_ocean_droplet: >
      state=present
      name=my_new_droplet
      client_id=XXX
      api_key=XXX
      size_id=1
      region_id=2
      image_id=3
      wait_timeout=500
  register: my_droplet
- debug: msg="ID is {{ my_droplet.droplet.id }}"
- debug: msg="IP is {{ my_droplet.droplet.ip_address }}"

# Ensure a droplet is present
# If droplet id already exist, will return the droplet details and changed = False
# If no droplet matches the id, a new droplet will be created and the droplet details (including the new id) are returned, changed = True.

- digital_ocean_droplet: >
      state=present
      id=123
      name=my_new_droplet
      client_id=XXX
      api_key=XXX
      size_id=1
      region_id=2
      image_id=3
      wait_timeout=500

# Create a droplet with ssh key
# The ssh key id can be passed as argument at the creation of a droplet (see ssh_key_ids).
# Several keys can be added to ssh_key_ids as id1,id2,id3
# The keys are used to connect as root to the droplet.

- digital_ocean_droplet: >
      state=present
      ssh_key_ids=id1,id2
      name=my_new_droplet
      client_id=XXX
      api_key=XXX
      size_id=1
      region_id=2
      image_id=3
'''

import sys
import os
import time

try:
    from dopy.manager import DoError, DoManager
except ImportError as e:
    print "failed=True msg='dopy required for this module'"
    sys.exit(1)

class TimeoutError(DoError):
    def __init__(self, msg, id):
        super(TimeoutError, self).__init__(msg)
        self.id = id

class JsonfyMixIn(object):
    def to_json(self):
        return self.__dict__

class Droplet(JsonfyMixIn):
    manager = None

    def __init__(self, droplet_json):
        self.status = 'new'
        self.__dict__.update(droplet_json)

    def is_powered_on(self):
        return self.status == 'active'

    def update_attr(self, attrs=None):
        if attrs:
            for k, v in attrs.iteritems():
                setattr(self, k, v)
        else:
            json = self.manager.show_droplet(self.id)
            if json['ip_address']:
                self.update_attr(json)

    def power_on(self):
        assert self.status == 'off', 'Can only power on a closed one.'
        json = self.manager.power_on_droplet(self.id)
        self.update_attr(json)

    def ensure_powered_on(self, wait=True, wait_timeout=300):
        if self.is_powered_on():
            return
        if self.status == 'off':  # powered off
            self.power_on()

        if wait:
            end_time = time.time() + wait_timeout
            while time.time() < end_time:
                time.sleep(min(20, end_time - time.time()))
                self.update_attr()
                if self.is_powered_on():
                    if not self.ip_address:
                        raise TimeoutError('No ip is found.', self.id)
                    return
            raise TimeoutError('Wait for droplet running timeout', self.id)

    def destroy(self):
        return self.manager.destroy_droplet(self.id)

    @classmethod
    def setup(cls, client_id, api_key):
        cls.manager = DoManager(client_id, api_key)

    @classmethod
    def add(cls, name, size_id, image_id, region_id, ssh_key_ids=None):
        json = cls.manager.new_droplet(name, size_id, image_id, region_id, ssh_key_ids)
        droplet = cls(json)
        return droplet

    @classmethod
    def find(cls, id=None, name=None):
        if not id and not name:
            return False

        droplets = cls.list_all()

        # Check first by id.  digital ocean requires that it be unique
        for droplet in droplets:
            if droplet.id == id:
                return droplet

        # Failing that, check by hostname.
        for droplet in droplets:
            if droplet.name == name:
                return droplet

        return False

    @classmethod
    def list_all(cls):
        json = cls.manager.all_active_droplets()
        return map(cls, json)

def core(module):
    def getkeyordie(k):
        v = module.params[k]
        if v is None:
            module.fail_json(msg='Unable to load %s' % k)
        return v

    try:
        # params['client_id'] will be None even if client_id is not passed in
        client_id = module.params['client_id'] or os.environ['DO_CLIENT_ID']
        api_key = module.params['api_key'] or os.environ['DO_API_KEY']
    except KeyError, e:
        module.fail_json(msg='Unable to load %s' % e.message)

    changed = True
    state = module.params['state']

    Droplet.setup(client_id, api_key)
    if state in ('present'):

        # First, try to find a droplet by id.
        droplet = Droplet.find(id=module.params['id'])

        # If we couldn't find the droplet and the user is allowing unique
        # hostnames, then check to see if a droplet with the specified
        # hostname already exists.
        if not droplet and module.params['unique_name']:
            droplet = Droplet.find(name=getkeyordie('name'))

        # If both of those attempts failed, then create a new droplet.
        if not droplet:
            droplet = Droplet.add(
                name=getkeyordie('name'),
                size_id=getkeyordie('size_id'),
                image_id=getkeyordie('image_id'),
                region_id=getkeyordie('region_id'),
                ssh_key_ids=module.params['ssh_key_ids']
            )

        if droplet.is_powered_on():
            changed = False

        droplet.ensure_powered_on(
            wait=getkeyordie('wait'),
            wait_timeout=getkeyordie('wait_timeout')
        )

        module.exit_json(changed=changed, droplet=droplet.to_json())

    elif state in ('absent'):
        # First, try to find a droplet by id.
        droplet = None
        if 'id' in module.params:
            droplet = Droplet.find(id=module.params['id'])

        # If we couldn't find the droplet and the user is allowing unique
        # hostnames, then check to see if a droplet with the specified
        # hostname already exists.
        if not droplet and module.params['unique_name'] and 'name' in module.params:
            droplet = Droplet.find(name=module.params['name'])

        if not droplet:
            module.exit_json(changed=False, msg='The droplet is not found.')

        event_json = droplet.destroy()
        module.exit_json(changed=True, event_id=event_json['event_id'])

def main():
    module = AnsibleModule(
        argument_spec = dict(
            state = dict(choices=['present', 'absent'], default='present'),
            client_id = dict(aliases=['CLIENT_ID'], no_log=True),
            api_key = dict(aliases=['API_KEY'], no_log=True),
            name = dict(type='str'),
            size_id = dict(type='int'),
            image_id = dict(type='int'),
            region_id = dict(type='int'),
            ssh_key_ids = dict(default=''),
            id = dict(aliases=['droplet_id'], type='int'),
            unique_name = dict(type='bool', choices=BOOLEANS, default='no'),
            wait = dict(type='bool', choices=BOOLEANS, default='yes'),
            wait_timeout = dict(default=300, type='int'),
        ),
        required_together = (
            ['size_id', 'image_id', 'region_id'],
        ),
        required_one_of = (
            ['id', 'name'],
        ),
    )

    try:
        core(module)
    except TimeoutError as e:
        module.fail_json(msg=str(e), id=e.id)
    except (DoError, Exception) as e:
        module.fail_json(msg=str(e))

# import module snippets
from ansible.module_utils.basic import *

main()