~launchpad-pqm/launchpad/devel

#!/usr/bin/python
#
# Copyright 2009 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

"""
Generate an EC2 image that is capable of running the Windmill browser UI
testing tool.

You must provide a base image that will be augmented with the necessary
packages and configuration.

The script requires certain options to be specified in order to function
properly.  These options may be supplied using command-line switches, or
via a config file, with the --config command-line switch.  The default
config file location is ~/.ec2/ec2bundle.cfg

The config file format simply replicates the required command-line options
as configuration keys.

---- ec2bundle.cfg ---

[DEFAULT]
key = gsg-keypair
identity-file = ~/.ec2/foo-keypair-id_rsa
private-key = ~/.ec2/pk-HKZYKTAIG2ECMXYIBH3HXV4ZBZQ55CLO.pem
cert =~/.ec2/cert-HKZYKTAIG2ECMXYIBH3HXV4ZBZQ55CLO.pem
user-id = AIDADH4IGTRXXKCD
access-key = SOMEBIGSTRINGOFDIGITS
secret-key = s0m3funKyStr1Ng0fD1gitZ
#bucket = foo  # Required, but you probably want to customize it each time.

---- fin ---

"""

__metatype__ = type


# Reuse a whole bunch of code from ec2test.py.
import ConfigParser
import ec2test
import logging
import optparse
import os
import paramiko
import select
import socket
import subprocess
import sys
import time


log   = logging.getLogger(__name__)
info  = log.info
debug = log.debug


usage = """
Generate an EC2 image for Windmill testing in Firefox.

usage: %prog [options] AMI-ID
"""

class Instance:
    """An EC2 instance controller."""

    def __init__(self, instance):
        self._instance = instance

    @property
    def id(self):
        return self._instance.id

    @property
    def hostname(self):
        return self._instance.public_dns_name

    def stop(self):
        instance = self._instance

        instance.update()
        if instance.state not in ('shutting-down', 'terminated'):
            # terminate instance
            instance.stop()
            instance.update()

        info('instance %s\n' % (instance.state,))

    def wait_for_instance_to_start(self):
        """Wait for the instance to transition to the "running" state."""
        instance = self._instance
        info('Instance %s starting..' % instance.id)

        start = time.time()
        while instance.state == 'pending':
            sys.stdout.write('.')
            sys.stdout.flush()
            time.sleep(5)
            instance.update()
        if instance.state == 'running':
            info('\ninstance now running at %s\n' % instance.public_dns_name)
            elapsed = time.time() - start
            info('Started in %d minutes %d seconds\n' %
                     (elapsed // 60, elapsed % 60))
            cout = instance.get_console_output()
            info(cout.output)
        else:
            raise RuntimeError('failed to start: %s\n' % instance.state)

    @classmethod
    def from_image(cls, account, ami_id, instance_type):
        """Return a new instance using the given startup parameters."""
        info("Starting instance")

        # Set up a security group that opens up ports 22, 80, and 443.  Also
        # opens up access for our IP.
        account.acquire_security_group()

        image = account.acquire_image(ami_id)
        key = account.name
        debug("Image: %s, Type: %s, Key: %s" % (
              ami_id, instance_type, key))

        reservation = image.run(
            key_name=key,
            security_groups=[key],
            instance_type=instance_type)

        instance = cls(reservation.instances[0])
        instance.wait_for_instance_to_start()
        return instance

    @classmethod
    def from_running_instance(cls, account, instance_id):
        """Create an object from an already running EC2 instance."""
        instance = account.get_instance(instance_id)
        if not instance:
            raise RuntimeError(
                "Unable to connect to instance %s" % instance_id)

        info("Connected to instance %s" % instance_id)
        proxy = cls(instance)
        # Just to be extra safe.
        proxy.wait_for_instance_to_start()
        return proxy


class SSHConnector:
    """Handle the various aspects of using an SSH connection."""

    def __init__(self, hostname, user, identity_file):
        self.hostname = hostname
        self.user = user
        self.identity_file = os.path.expanduser(identity_file)
        self._client = None

    def get_private_key(self):
        """Generate a private key object for our keyfile"""
        fp = os.path.expanduser(self.identity_file)
        return paramiko.RSAKey.from_private_key(open(fp))

    def connect(self):
        info('Waiting for SSH to come available: %s@%s\n' % (
             self.user, self.hostname))
        debug("Using private key file: %s" % self.identity_file)

        private_key = self.get_private_key()

        for count in range(10):
            self._client = paramiko.SSHClient()
            self._client.set_missing_host_key_policy(ec2test.AcceptAllPolicy())

            try:

                self._client.connect(
                    self.hostname,
                    username=self.user,
                    pkey=private_key,
                    allow_agent=False,
                    look_for_keys=False)

            except (socket.error, paramiko.AuthenticationException), e:
                log.warning('wait_for_connection: %r' % (e,))
                if count < 9:
                    time.sleep(5)
                    info('retrying...')
                else:
                    raise
            else:
                break

    def exec_command(self, remote_command, check_return=True):
        """Execute a command on the remote server.

        Raises an error if the command returns an exit status that is not
        zero, unless the option `check_return=False' has been given.
        """
        info('Executing command: %s@%s %s\n' % (self.user, self.hostname, remote_command))

        session = self._client.get_transport().open_session()
        session.exec_command(remote_command)
        session.shutdown_write()

        # TODO: change this to use the logging module
        while True:
            select.select([session], [], [], 0.5)
            if session.recv_ready():
                data = session.recv(4096)
                if data:
                    sys.stdout.write(data)
                    sys.stdout.flush()
            if session.recv_stderr_ready():
                data = session.recv_stderr(4096)
                if data:
                    sys.stderr.write(data)
                    sys.stderr.flush()
            if session.exit_status_ready():
                break
        session.close()

        # XXX: JonathanLange 2009-05-31: If the command is killed by a signal
        # on the remote server, the SSH protocol does not send an exit_status,
        # it instead sends a different message with the number of the signal
        # that killed the process. AIUI, this code will fail confusingly if
        # that happens.
        exit_status = session.recv_exit_status()
        if exit_status and check_return:
            raise RuntimeError('Command failed: %s' % (remote_command,))
        return exit_status

    def copy_to_remote(self, local_filename, remote_filename):
        cmd = [
            'scp',
            '-i', self.identity_file,
            local_filename,
            '%s@%s:%s' % (self.user, self.hostname, remote_filename)
            ]
        info("Executing command: %s" % ' '.join(cmd))
        subprocess.check_call(cmd)

    def user_command(self):
        """Return a user-friendly ssh command-line string."""
        return "ssh -i %s %s@%s" % (
            self.identity_file,
            self.user,
            self.hostname)


class ImageBundler:
    """Bundle an EC2 image on a remote system."""

    def __init__(self, private_key, cert, account_id, target_bucket,
                 access_key, secret_key, ssh):
        self.private_key = os.path.expanduser(private_key)
        self.cert = os.path.expanduser(cert)
        self.account_id = account_id
        self.target_bucket = target_bucket
        self.access_key = access_key
        self.secret_key = secret_key
        self.ssh = ssh

        # Use the instance /mnt directory by default, because it has a few
        # hundred GB of free space to work with.
        self._bundle_dir = os.path.join('/mnt', target_bucket)

    def bundle_image(self):
        self.configure_bundling_environment()
        manifest = self._bundle_image()
        self._upload_bundle(manifest)
        self._register_image(manifest)

    def remote_private_keypath(self):
        # ALWAYS have these files in /mnt on the remote system.  Otherwise
        # they will get bundled along with the image.
        return os.path.join('/mnt', os.path.basename(self.private_key))

    def remote_certpath(self):
        # ALWAYS have these files in /mnt on the remote system.  Otherwise
        # they will get bundled along with the image.
        return os.path.join('/mnt', os.path.basename(self.cert))

    def configure_bundling_environment(self):
        """Configure what we need on the instance for bundling the image."""
        # Send our keypair to the remote environment so that it can be used
        # to bundle the image.
        local_cert = os.path.abspath(self.cert)
        local_pkey = os.path.abspath(self.private_key)

        # ALWAYS copy these files into /mnt on the remote system.  Otherwise
        # they will get bundled along with the image.
        remote_cert = self.remote_certpath()
        remote_pkey = self.remote_private_keypath()

        # See if the files are present, and copy them over if they are not.
        self._ensure_remote_file(remote_cert, local_cert)
        self._ensure_remote_file(remote_pkey, local_pkey)

    def _ensure_remote_file(self, remote_file, desired_file):
        info("Checking for '%s' on the remote system" % remote_file)
        test = 'test -f %s' % remote_file
        exit_status = self.ssh.exec_command(test, check_return=False)
        if bool(exit_status):
            self.ssh.copy_to_remote(desired_file, remote_file)

    def _bundle_image(self):
        # Create the bundle in a subdirectory, to avoid spamming up /mnt.
        self.ssh.exec_command(
            'mkdir %s' % self._bundle_dir, check_return=False)

        cmd = [
            'ec2-bundle-vol',
            '-d %s' % self._bundle_dir,
            '-b',   # Set batch-mode, which doesn't use prompts.
            '-k %s' % self.remote_private_keypath(),
            '-c %s' % self.remote_certpath(),
            '-u %s'  % self.account_id
            ]

        self.ssh.exec_command(' '.join(cmd))
        # Assume that the manifest is 'image.manifest.xml', since "image" is
        # the default prefix.
        manifest = os.path.join(self._bundle_dir, 'image.manifest.xml')

        # Best check that the manifest actually exists though.
        test = 'test -f %s' % manifest
        exit_status = self.ssh.exec_command(test, check_return=False)

        if bool(exit_status):
            raise RuntimeError(
                "Failed to write the image manifest file: %s" % manifest)

        return manifest

    def _upload_bundle(self, manifest):
        cmd = [
            'ec2-upload-bundle',
            '-b %s' % self.target_bucket,
            '-m %s' % manifest,
            '-a %s' % self.access_key,
            '-s %s' % self.secret_key
            ]
        self.ssh.exec_command(' '.join(cmd))

    def _register_image(self, manifest):
        # This is invoked locally.
        mfilename = os.path.basename(manifest)
        manifest_path = os.path.join(self.target_bucket, mfilename)

        env = os.environ.copy()
        env['JAVA_HOME'] = '/usr/lib/jvm/default-java'
        cmd = [
            'ec2-register',
            '--private-key=%s' % self.private_key,
            '--cert=%s' % self.cert,
            manifest_path
            ]
        info("Executing command: %s" % ' '.join(cmd))
        subprocess.check_call(cmd, env=env)


class XvfbSystemConfigurator:
    """Configure a remote operating system over SSH to use the xvfb server."""

    def __init__(self, ssh):
        self.ssh = ssh

    def configure_system(self):
        """Configure the operating system with the needed packages, etc."""
        do = self.ssh.exec_command

        # Make sure we know about all the packages, and where they may be
        # found.
        do("apt-get -y update")
        # Install the necessary packages
        do("apt-get -y install xvfb firefox xfonts-base")


class CombinedConfigParser:
    """Store and reconcile options for both optparse and ConfigParser."""

    def __init__(self, optparser, cfgparser):
        self._optparser = optparser
        self._cfgparser = cfgparser

        # A list of required optparse options.
        self.required_options = []
        self.known_options = []

        # Our parsed positional command-line arguments, as returned by
        # optparse.OptionParser.parse_args()
        self.args = None

        # An optparse option.dest to 'cfg-key' mapping.
        self._option_to_cfgkey = {}

        self._parsed_cli_options = None
        self._parsed_cfg_options = None

    def __getattr__(self, name):
        return self.get(name)

    def add_option(self, *args, **kwds):
        """Wrap the OptionParser.add_option() method, and add our options."""
        try:
            # We can't pass unknown kwds to make_option, or it will barf.
            is_required = kwds.pop('required')
        except KeyError:
            is_required = False

        option = optparse.make_option(*args, **kwds)
        self._optparser.add_option(option)

        if is_required:
            self.add_required_option(option)

        self._add_option_to_cfg_mapping(option, args)

    def add_required_option(self, option):
        """Add a required option.

        Takes an optparse.Option object.
        """
        self.required_options.append(option)

    def _add_option_to_cfg_mapping(self, option, option_constructor_args):
        # Convert the long options into .ini keys.  Use the last long option
        # given.
        for switch in reversed(option_constructor_args):
            if switch.startswith('--'):
                # We found a '--foo' switch, so use it.  Drop the '--',
                # because the config file doesn't use the prefixes.
                self._option_to_cfgkey[option.dest] = switch[2:]

    def error(self, message):
        """Wrap optparse.OptionParser.error()."""
        self._optparser.error(message)

    def parse_config_file(self, filepath):
        fp = os.path.expanduser(filepath)

        if not os.path.exists(fp):
            self.error("The config file '%s' does not exist!" % fp)

        self._cfgparser.read(fp)
        self._parsed_cfg_options = self._cfgparser.defaults()

        num_opts = len(self._parsed_cfg_options)
        debug("Loaded %d options from %s" % (num_opts, fp))

    def parse_cli_args(self, argv):
        """Wrap optparse.OptionParser.parse_args()."""
        options, args = self._optparser.parse_args(argv)
        self._parsed_cli_options = options
        self.args = args
        return (options, args)

    def verify_options(self):
        """Verify that all required options are there.

        Raise an optparse.OptionParser.error() if something is missing.

        Make sure you parsed the config file with parse_config_file() before
        doing this.
        """
        debug("Verifying options")
        if not self._parsed_cfg_options:
            debug("No config file options found")

        for option in self.required_options:
            # Check for a command-line option.

            option_name = option.dest

            if self.get(option_name) is None:
                self._required_option_error(option)
            else:
                debug("Found required option: %s" % option_name)

    def _required_option_error(self, option):
        msg = "Required option '%s' was not given (-h for help)" % str(option)
        self.error(msg)

    def get(self, name, default=None):
        """Return the appropriate option, CLI first, CFG second."""
        cli_name = name
        cfg_name = self._option_to_cfgkey.get(name)

        value = self._getoption(cli_name)

        if value is None and cfg_name is not None:
            # No command-line option was supplied, but we do have a config
            # file entry with that name.
            value = self._getcfg(cfg_name)

            if value is None:
                # No config file option was supplied either, so return the
                # default.
                return default

        return value

    def _getoption(self, key, default=None):
        return getattr(self._parsed_cli_options, key, default)

    def _getcfg(self, key, default=None):
        return self._parsed_cfg_options.get(key, default)


def get_credentials():
    """Return an EC2Credentials object for accessing the webservice."""
    # Get the AWS identifier and secret identifier.
    return ec2test.EC2Credentials.load_from_file()


def parse_config_file(filepath):
    config = ConfigParser.ConfigParser()
    config.read(filepath)
    return config


def parse_options(argv):
    oparser = optparse.OptionParser(usage)
    cparser = ConfigParser.SafeConfigParser()
    parser = CombinedConfigParser(oparser, cparser)

    # What follows are "Required options" - these must be supplied from either
    # the command-line, or from a config file.
    parser.add_option(
        '-k', '--key',
        dest="keypair_name",
        required=True,
        help="The name of the AWS key pair to use for launching instances.")

    parser.add_option(
        '-K', '--private-key',
        dest="private_key",
        required=True,
        help="The X.509 private keyfile that will be used to sign the new "
             "image.")

    parser.add_option(
        '-C', '--cert',
        dest="cert",
        required=True,
        help="The X.509 certificate that will be used to bundle the new "
             "image.")

    parser.add_option(
        '-i', '--identity-file',
        dest='identity_file',
        required=True,
        help="The location of the RSA private key that SSH will use to "
             "connect to the instance.")

    parser.add_option(
        '-b', '--bucket',
        dest="bucket",
        required=True,
        help="The bucket that the image will be placed into.")

    parser.add_option(
        '-u', '--user-id',
        dest="account_id",
        required=True,
        help="Your 12 digit AWS account ID")

    parser.add_option(
        '-a', '--access-key',
        dest="access_key",
        required=True,
        help="Your AWS access key.")

    parser.add_option(
        '-s', '--secret-key',
        dest="secret_key",
        required=True,
        help="Your AWS secret key.")


    # Start our "Optional options."
    parser.add_option(
        '-v', '--verbose',
        action='store_true',
        dest='verbose',
        default=False,
        help="Turn on debug output.")

    parser.add_option(
        '-t', '--instance-type',
        dest="instance_type",
        default="m1.large",
        help="The type of instance to be launched.  Should be the same as "
             "the base image's required type. [default: %default]")

    parser.add_option(
        '-c', '--config',
        dest='config',
        default="~/.ec2/ec2bundle.cfg",
        help="Load script options from the supplied config file. (.ini "
             "format, see the module docstring for details.) "
             "[default: %default]")

    parser.add_option(
        '--keepalive',
        action='store_true',
        dest="keepalive",
        default=False,
        help="Don't shut down the instance when we are done building (or "
             "erroring out).")

    parser.add_option(
        '--no-bundle',
        action='store_true',
        dest='no_bundle',
        default=False,
        help="Don't create a bundle, just start the server and configure the "
             "environment.")

    parser.add_option(
        '--use-instance',
        dest='running_instance',
        help="Use the supplied EC2 instance ID, instead of starting our own "
             "server.  The instance will be left running.")


    options, args = parser.parse_cli_args(argv)

    # Do this ASAP
    if options.verbose:
        log.setLevel(logging.DEBUG)

    if options.config:
        parser.parse_config_file(options.config)

    # Make sure all the required args are present.  Will error-out if
    # something is missing.
    parser.verify_options()

    if len(args) != 2:
        parser.error("You must provide an AMI ID that can serve as the new "
                     "image's base.")

    return parser


def main(argv):
    config = parse_options(argv)

    credentials = get_credentials()
    account = credentials.connect(config.keypair_name)

    # Save the flag so we can change it.  This is how we enforce shutdown
    # policies.
    keepalive = config.keepalive

    if config.running_instance:
        # We want to keep the server alive if the user supplied their own
        # instance.  Killing it without their consent would be cruel.
        keepalive = True

    ssh_user_command = None
    try:
        try:
            instance = None
            if config.running_instance:
                # Connect to an already running instance.
                instance = Instance.from_running_instance(
                    account, config.running_instance)
            else:
                # Run an instance for our base image.
                instance = Instance.from_image(
                    account, config.args[1], config.instance_type)

            ssh = SSHConnector(
                instance.hostname, 'root', config.identity_file)
            ssh.connect()
            ssh_user_command = ssh.user_command()

            system_configurator = XvfbSystemConfigurator(ssh)
            system_configurator.configure_system()

            if not config.no_bundle:
                bundler = ImageBundler(
                    config.private_key,
                    config.cert,
                    config.account_id,
                    config.bucket,
                    config.access_key,
                    config.secret_key,
                    ssh)
                bundler.bundle_image()

        except:
            # Log the exception now so it doesn't interfere with or get eaten
            # by the instance shutdown.
            log.exception("Oops!")
    finally:
        if keepalive:
            log.warning("instance %s is now running on its own" % instance.id)
            if ssh_user_command:
                info("You may now ssh into the instance using the following command:")
                info("  $ %s" % ssh_user_command)

            log.warning("Remember to shut the instance down when you are done!")
        else:
            instance.stop()


if __name__ == '__main__':
    logging.basicConfig()
    main(sys.argv)