euca-install-load-balancer

#!/usr/bin/env python
#
# Software License Agreement (BSD License)
#
# Copyright (c) 2013, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software in source and binary forms, with or
# without modification, are permitted provided that the following conditions
# are met:
#
#   Redistributions of source code must retain the above
#   copyright notice, this list of conditions and the
#   following disclaimer.
#
#   Redistributions in binary form must reproduce the above
#   copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other
#   materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.


import argparse
import copy
import glob
import logging
import os
import re
import subprocess
import sys
import boto
from requestbuilder.exceptions import ServiceInitError
import euca2ools
import time

(major, minor, patch) = euca2ools.__version__.split('-')[0].split('.')
if int(major) < 3 or (int(major) >= 3 and int(minor) < 1):
    print >> sys.stderr, "euca2ools version 3.1.0 or newer required."
    sys.exit(1)

DEFAULT_IMAGE_LOCATION = '/usr/share/eucalyptus-load-balancer-image'

class LBManager(object):
    BUCKET_PREFIX = 'loadbalancer'
    IMAGE_RE = re.compile('eucalyptus-load-balancer-image')

    def __init__(self):
        self._check_environment()
        self._populate_images()

    def _check_environment(self):
        env = self._get_env()
        if not "EC2_URL" in env:
            print >> sys.stderr, "Error: Unable to find EC2_URL"
            print >> sys.stderr, "Make sure your eucarc is sourced."
            sys.exit(1)
        cmd = ['euca-modify-property']
        try:
            with open(os.devnull, 'w') as nullfile:
                subprocess.call(cmd, env=env, stdout=nullfile)
        except OSError:
            print >> sys.stderr, "Error: cannot find 'euca-modify-property' binary."
            print >> sys.stderr, "Make sure EUCALYPTUS path variable is exported."
            sys.exit(1)

    def _populate_images(self):
        self.images = {}
        for image in boto.connect_ec2_endpoint(os.environ.get("EC2_URL")).get_all_images():
            if self.IMAGE_RE.search(image.location):
                version = self._get_image_version(image)
                if not version in self.images:
                    self.images[version] = []
                self.images[version].append(image)

    def _get_image_version(self, image):
        manifest = image.location
        ### This looks for - or _ because 3.4.x used buckets named with _
        matches = re.match(r'{0}[_|-]v(\d+)'.format(self.BUCKET_PREFIX), manifest)
        try:
            return int(matches.group(1))
        except:
            return 0

    def _build_bucket_name(self):
        return "{0}-v{1}".format(self.BUCKET_PREFIX,
                                 self._get_next_version())

    def _split_location(self, location):
        matches = re.match(r'(.+)/(.+)\.manifest\.xml', location)
        return matches.groups()

    def _get_env(self):
        return os.environ.copy()

    def _get_next_version(self):
        if not self.images:
            return 1
        else:
            return sorted(self.images)[-1] + 1

    def _remove(self, images, force=False):
        removed = []
        enabled_image = self.get_enabled()
        for image_set in images.itervalues():
            should_remove = True
            if not force:
                for image in image_set:
                    imageid = image.id
                    if imageid == enabled_image:
                        print >> sys.stderr, "Warning: skipping enabled image."
                        print >> sys.stderr, "Use '--force' to remove it anyway."
                        should_remove = False

            if should_remove:
                for image in image_set:
                    imageid = image.id
                    deregister_cmd = "euca-deregister {0}".format(imageid)
                    deregister_output = subprocess.Popen(deregister_cmd, stdout=subprocess.PIPE,
                                                            stderr=subprocess.PIPE, shell=True)
                    deregister_output.wait()
                    location = image.location
                    (bucket, prefix) = self._split_location(location)
                    delete_bundle_cmd = "euca-delete-bundle -b {0}/{1}".format(bucket, prefix)
                    delete_bundle_output = subprocess.Popen(delete_bundle_cmd, stdout=subprocess.PIPE,
                                                            stderr=subprocess.PIPE, shell=True)
                    delete_bundle_output.wait()
                    removed.append(imageid)
        return removed

    def _run_command(self, command):
        pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
        pipe.wait()
        stdout = pipe.stdout.read()
        stderr = pipe.stderr.read()
        if pipe.returncode > 0:
            print "Error: Unable to install load balancer image due to:\n" \
                                 + stdout + stderr
            sys.exit(1)
        return (stdout, stderr)

    def remove_all(self, force=False):
        return self._remove(self.images, force=force)

    def remove_old(self, force=False):
        old_images = copy.deepcopy(self.images)
        newest_key = sorted(old_images)[-1]
        del old_images[newest_key]
        return self._remove(old_images, force=force)

    def list_images(self):
        if not self.images:
            print "No load balancer images Installed."
            print
            return

        enabled_image = self.get_enabled()
        for version, images in self.images.iteritems():
            lines = []
            enabled_marker = False
            for image in images:
                if enabled_image == image.id:
                    enabled_marker = True
                lines.append("{0} ({1})".format(image.id,
                                                image.location))
                if image.description is not None:
                    lines.append('\t{0}'.format(image.description))
            if enabled_marker:
                print "Version {0} (enabled)".format(version)
            else:
                print "Version {0}".format(version)
            for line in lines:
                print line
            print

    def install_and_enable(self, tarball):
        image = self.install(tarball)
        if image is None:
            print >> sys.stderr, "Error: Load balancer installation failed."
            sys.exit(1)
        self.enable(image)
        print
        print "Load Balancing Support is Enabled"

    def install(self, tarball):
        ### Decompress image
        decompress_cmd = "/bin/tar xzfv {0}".format(tarball)
        print "Decompressing tarball: " + tarball
        decompress_stdout, decompress_stderr = self._run_command(decompress_cmd)
        ### Bundle and upload image
        bucket = self._build_bucket_name()
        image_file = decompress_stdout.strip()
        install_cmd = "euca-bundle-and-upload-image -b {0} -i {1} -r x86_64".format(bucket, image_file)
        print "Bundling and uploading image to bucket: " + bucket
        bundle_stdout, bundle_stderr = self._run_command(install_cmd)
        try:
            manifest = bundle_stdout.split()[1]
        except IndexError:
            print "Error: Unable to retrieve uploaded image manifest: " + manifest
            sys.exit(1)

        register_cmd = "euca-register {0} --name {1} -d '{2}' " \
                       "--virtualization-type hvm".format(manifest, "eucalyptus-load-balancer-image-v"
                        + str(self._get_next_version()), time.strftime(
                                'Installed on %Y-%m-%d at %H:%M:%S %Z'))
        print "Registering image manifest: " + manifest
        register_stdout, register_stderr = self._run_command(register_cmd)
        emi_id = register_stdout.split()[1]
        print "Registered image: " + emi_id

        remove_tarball_cmd = "rm -f {0}".format(image_file)
        remove_stdout, remove_stderr = self._run_command(remove_tarball_cmd)
        return emi_id

    def get_enabled(self):
        try:
            env = self._get_env()
            cmd = ['euca-describe-properties', 'loadbalancing.loadbalancer_emi']
            out = subprocess.Popen(cmd, env=env,
                                   stdout=subprocess.PIPE).communicate()[0]
            imageid = out.split()[-1]
            if imageid == "NULL":
                return None
            else:
                return imageid
        except OSError:
            print >> sys.stderr, "Error: failed to get Load Balancer EMI."
            sys.exit(1)

    def enable(self, imageid):
        try:
            env = self._get_env()
            cmd = ['euca-modify-property', '-p',
                   'loadbalancing.loadbalancer_emi={0}'.format(imageid)]
            subprocess.check_call(cmd, env=env)
        except (OSError, subprocess.CalledProcessError):
            print >> sys.stderr, "Error: failed to set Load Balancer EMI."
            print >> sys.stderr, "You'll have to enable it manually."
            print >> sys.stderr
            print >> sys.stderr, "To enable load balancing service support, run this command:"
            print >> sys.stderr, " ".join(cmd)
            sys.exit(1)

    @classmethod
    def create(cls):
        #
        # Quiet the attempts to get logger for Walrus
        #
        class NullHandler(logging.Handler):
            def emit(self, record):
                pass
        logging.getLogger("Walrus").addHandler(NullHandler())

        try:
            return cls()
        except ServiceInitError as err:
            print >> sys.stderr, str(err)
            sys.exit(1)
        

if __name__ == "__main__":
    description = '''
    Load Balancer Installation Tool:

    This tool provides an easy way to install a Eucalyptus Load
    Balancer image.  Normally, you'll want to pass '--default-tarball'
    to install the latest Load Balancer image and register it with
    Eucalyptus. You may next want to pass '--remove-old' if you need
    to clean up older versions of the Load Balancer.

    Installed Load Balancers will be marked with a "version" number
    that differentiates one installed Load Balancer from another. The
    larger the "version", the more recently the Load Balancer has been
    installed.

    WARNING: DO NOT ATTEMPT TO REMOVE LOAD BALANCER IMAGES WHILE
    INSTANCES OF THAT LOAD BALANCER ARE STILL RUNNING. FIRST TERMINATE
    ANY RUNNING LOAD BALANCERS BEFORE ATTEMPTING THIS!'''

    epilog = '''
    NOTE: In order to use this you MUST have cloud administrator
    credentials sourced in your environment (i.e., run the command
    '. /my/cred/path/eucarc').'''

    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=description, epilog=epilog)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-t', '--tarball', metavar='TARBALL',
                       help='Load Balancer tarball to install')
    group.add_argument('--install-default', action='store_true',
                       help='''This option must be supplied if you
                       would like to install the default tarball''')
    group.add_argument('-l', '--list', action='store_true',
                       help='''List currently installed Load Balancer
                       image bundles.''')
    group.add_argument('--remove-old', action='store_true',
                       help='''Remove OLD Load Balancers. Your most
                       recently installed Load Balancer and your
                       currently enabled Load Balancer will not be
                       removed. Most of the time, these will be the
                       same Load Balancer.''')
    group.add_argument('--remove-all', action='store_true',
                       help='''Remove ALL Load Balancers. The
                       currently enabled Load Balancer will be skipped
                       unless '--force' is passed as well.''')
    parser.add_argument('--force', action='store_true', help='''Force
                        an operation. This will force removal of
                        enabled Load Balancers.''')

    args = parser.parse_args()
    lbm = LBManager.create()

    if args.tarball:
        lbm.install_and_enable(args.tarball)
    elif args.install_default:
        print 'Installing default Load Balancer tarball.'
        try:
            tarball = glob.glob('{0}*/*.tgz'.format(DEFAULT_IMAGE_LOCATION))[0]
            print 'Found tarball {0}'.format(tarball)
            lbm.install_and_enable(tarball)
        except IndexError:
            print >> sys.stderr, "Error: failed to find a Load Balancer tarball."
            print >> sys.stderr, "Try supplying one on the command line with '-t'."
            sys.exit(1)
    elif args.remove_old:
        imageids = lbm.remove_old(force=args.force)
        for imageid in imageids:
            print "Removed {0}".format(imageid)
    elif args.remove_all:
        imageids = lbm.remove_all(force=args.force)
        for imageid in imageids:
            print "Removed {0}".format(imageid)
    elif args.list:
        print "Currently Installed Load Balancer Bundles:"
        print
        lbm.list_images()