Draft of changes to salt-cloud EC2.py

EC2.py

'''
The generic libcloud template used to create the connections and deploy the
cloud virtual machines
'''

# Import python libs
import os
import sys
import subprocess

# Import libcloud
from libcloud.compute.types import Provider
from libcloud.compute.providers import get_driver
from libcloud.compute.deployment import MultiStepDeployment, ScriptDeployment, SSHKeyDeployment

# Import salt libs
import saltcloud.utils

# Import paramiko
import paramiko


EC2_LOCATIONS = {
    'ap-northeast-1': Provider.EC2_AP_NORTHEAST,
    'ap-southeast-1': Provider.EC2_AP_SOUTHEAST,
    'eu-west-1': Provider.EC2_EU_WEST,
    'sa-east-1': Provider.EC2_SA_EAST,
    'us-east-1': Provider.EC2_US_EAST,
    'us-west-1': Provider.EC2_US_WEST,
    'us-west-2': Provider.EC2_US_WEST_OREGON,
}
DEFAULT_LOCATION = 'us-east-1'


def get_conn(location):
    '''
    Return a conn object for the passed vm data
    '''
    if location not in EC2_LOCATIONS:
        return None
    driver = get_driver(EC2_LOCATIONS[location])
    return driver(
            __opts__['EC2.id'],
            __opts__['EC2.key'],
            )


def ssh_pub(vm_):
    '''
    Deploy the primary ssh authentication key
    '''
    ssh = ''
    if 'ssh_auth' in vm_:
        if not os.path.isfile(vm_['ssh_auth']):
            return None
        ssh = vm_['ssh_auth']
    if not ssh:
        if not os.path.isfile(__opts__['ssh_auth']):
            return None
        ssh = __opts__['ssh_auth']

    return SSHKeyDeployment(open(os.path.expanduser(ssh)).read())


def script(vm_):
    '''
    Return the script deployment object
    '''
    minion = saltcloud.utils.minion_conf_string(__opts__, vm_)
    return ScriptDeployment(
            saltcloud.utils.os_script(
                saltcloud.utils.get_option(
                    'os',
                    __opts__,
                    vm_
                    ),
                vm_,
                __opts__,
                minion,
                ),
            name='/home/ec2-user/deployment.sh'
            )


def avail_images(location=None):
    '''
    Return a dict of all available vm images on the cloud provider with
    relevant data
    '''
    conn = get_conn(location)
    images = conn.list_images()
    ret = {}
    for img in images:
        ret[img.name] = {}
        for attr in dir(img):
            if attr.startswith('_'):
                continue
            ret[img.name][attr] = getattr(img, attr)
    return ret


def avail_sizes(location=None):
    '''
    Return a dict of all available vm images on the cloud provider with
    relevant data
    '''
    conn = get_conn(location)
    sizes = conn.list_sizes()
    ret = {}
    for size in sizes:
        ret[size.name] = {}
        for attr in dir(size):
            if attr.startswith('_'):
                continue
            ret[size.name][attr] = getattr(size, attr)
    return ret


def get_image(conn, vm_):
    '''
    Return the image object to use
    '''
    images = conn.list_images()
    if not 'image' in vm_:
        return images[0]
    if isinstance(vm_['image'], int):
        return images[vm_['image']]
    for img in images:
        if img.id == vm_['image']:
            return img


def get_size(conn, vm_):
    '''
    Return the vm's size object
    '''
    sizes = conn.list_sizes()
    if not 'size' in vm_:
        return sizes[0]
    if isinstance(vm_['size'], int):
        return sizes[vm_['size']]
    for size in sizes:
        if size.id == vm_['size']:
            return size
        if size.name == vm_['size']:
            return size


def get_location(vm_):
    '''
    Return the AWS region to use
    '''
    return vm_.get('location', __opts__.get('EC2.location', DEFAULT_LOCATION))


def get_availability_zone(conn, vm_):
    '''
    Return the availability zone to use
    '''
    zones = conn.list_locations()
    az = None
    if 'availability_zone' in vm_:
        az = vm_['availability_zone']
    elif 'EC2.availability_zone' in __opts__:
        az = __opts__['EC2.availability_zone']
    if az is None:
        # Default to first zone
        return zones[0].availability_zone
    for zone in zones:
        if zone.availability_zone == az:
            return az


def ssh_username(vm_):
    '''
    Return the ssh_username
    '''
    return vm_.get('ssh_username', __opts__.get('EC2.ssh_username', 'ec2-user'))


def keyname(vm_):
    '''
    Return the keyname
    '''
    return vm_.get('EC2.keyname', __opts__.get('EC2.keyname', ''))


def securitygroup(vm_):
    '''
    Return the securitygroup
    '''
    return vm_.get(
            'EC2.securitygroup',
            __opts__.get('EC2.securitygroup', 'default')
            )


def list_nodes():
    '''
    Return a list of the vms that are on the provider
    '''
    conn = get_conn() 
    nodes = conn.list_nodes()
    ret = {}
    for node in nodes:
        ret[node.name] = {
                'id': node.id,
                'image': node.image,
                'private_ips': node.private_ips,
                'public_ips': node.public_ips,
                'size': node.size,
                'state': node.state}
    return ret


def create(vm_):
    '''
    Create a single vm from a data dict
    '''
    location = get_location(vm_)
    print('Creating Cloud VM {0} in {1}'.format(vm_['name'], location))
    conn = get_conn(location)
    kwargs = {'ssh_username': ssh_username(vm_),
              'ssh_key': __opts__['EC2.private_key']}
    kwargs['name'] = vm_['name']
    kwargs['deploy'] = script(vm_)
    kwargs['image'] = get_image(conn, vm_)
    kwargs['size'] = get_size(conn, vm_)
    ex_keyname = keyname(vm_)
    if ex_keyname:
        kwargs['ex_keyname'] = ex_keyname
    ex_securitygroup = securitygroup(vm_)
    if ex_securitygroup:
        kwargs['ex_securitygroup'] = ex_securitygroup
    try:
        data = conn.deploy_node(**kwargs)
    except Exception as exc:
        err = ('The following exception was thrown by libcloud when trying to '
               'run the initial deployment: \n{0}\n\nThe vm {1} has been '
               'created but Salt could not be intsalled. Please verify that '
               'your ssh keys are in order and that the security group is '
               'accepting inbound connections from port 22.\n').format(
                       exc, vm_['name']
                       )
        sys.stderr.write(err)
        return False
    cmd = ('ssh -oStrictHostKeyChecking=no -t -i {0} {1}@{2} "sudo '
           '/home/{3}/deployment.sh"').format(
                   __opts__['EC2.private_key'],
                   kwargs['ssh_username'],
                   data.public_ips[0],
                   kwargs['ssh_username']
                   )
    subprocess.call(cmd, shell=True)
    print('Created Cloud VM {0} with the following values:'.format(
        vm_['name']
        ))
    for key, val in data.__dict__.items():
        print('  {0}: {1}'.format(key, val))