Charm: oneiric/buildbot-slave   Revision: 8   Hook: upgrade-charm
#!/usr/bin/env python

# Copyright 2012 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

import os
import shlex
import subprocess
import sys
import tempfile


def run(*args, **kwargs):
    """Run the command with the given arguments.

    The first argument is the path to the command to run.
    Subsequent arguments are command-line arguments to be passed.

    This function accepts all optional keyword arguments accepted by
    `subprocess.Popen`.
    """
    args = [i for i in args if i is not None]
    pipe = subprocess.PIPE
    process = subprocess.Popen(
        args, stdout=kwargs.pop('stdout', pipe),
        stderr=kwargs.pop('stderr', pipe),
        close_fds=kwargs.pop('close_fds', True), **kwargs)
    stdout, stderr = process.communicate()
    if process.returncode:
        raise subprocess.CalledProcessError(
            process.returncode, repr(args), output=stdout+stderr)
    return stdout


def command(*base_args):
    """Return a callable that will run the given command with any arguments.

    The first argument is the path to the command to run, subsequent arguments
    are command-line arguments to "bake into" the returned callable.

    The callable runs the given executable and also takes arguments that will
    be appeneded to the "baked in" arguments.

    For example, this code will list a file named "foo" (if it exists):

        ls_foo = command('/bin/ls', 'foo')
        ls_foo()

    While this invocation will list "foo" and "bar" (assuming they exist):

        ls_foo('bar')
    """
    def callable_command(*args):
        all_args = base_args + args
        return run(*all_args)

    return callable_command


log = command('juju-log')
apt_get_install = command('apt-get', 'install', '-y', '--force-yes')


def install_extra_repository(extra_repository):
    distribution = run('lsb_release', '-cs').strip()
    # Starting from Oneiric, `apt-add-repository` is interactive by
    # default, and requires a "-y" flag to be set.
    assume_yes = None if distribution == 'lucid' else '-y'
    try:
        run('apt-add-repository', assume_yes, extra_repository)
        run('apt-get', 'update')
    except subprocess.CalledProcessError as e:
        log('Error adding repository: ' + extra_repository)
        log(str(e))
        raise


def install_packages():
    install_extra_repository('ppa:yellow/ppa')
    apt_get_install('python-shell-toolbox')


install_packages()

# helpers and local depend on shelltoolbox so they cannot be imported until
# after the python-shell-toolbox package is installed.
from helpers import (
    get_config,
    log_entry,
    log_exit,
    )
from local import (
    config_json,
    create_slave,
    )

bzr = command('bzr')


def bzr_fetch(source, path):
    apt_get_install('bzr')
    target = tempfile.mktemp()
    bzr('branch', source, target)
    return os.path.join(target, path)


def bzr_cat(source, path):
    apt_get_install('bzr')
    content = bzr('--no-plugins', 'cat', source)
    target = os.path.join('/tmp', path)
    with open(target, 'w') as f:
        f.write(content)
    return target


def wget(source, path):
    target = os.path.join('/tmp', path)
    run('wget', '-O', target, source)
    return target


def hg_fetch(source, path):
    apt_get_install('mercurial')
    target = tempfile.mktemp()
    run('hg', 'clone', source, target)
    return os.path.join(target, path)


def git_fetch(source, path):
    apt_get_install('git')
    target = tempfile.mktemp()
    run('git', 'clone', source, target)
    return os.path.join(target, path)


METHODS = {
    'bzr': bzr_fetch,
    'bzr_cat': bzr_cat,
    'wget': wget,
    'hg': hg_fetch,
    'git': git_fetch,
    }


def handle_script(retrieve, url, path, args):
    log('Retrieving script: {}.'.format(url))
    script = retrieve(url, path)
    log('Changing script mode.')
    run('chmod', '+x', script)
    log('Executing script: {}.'.format(script))
    return subprocess.call([script] + shlex.split(str(args)))


def main():
    config = get_config()
    method = config.get('script-retrieval-method')
    url = config.get('script-url')
    path = config.get('script-path')
    # This is a naive substitution.  We can make it more sophisticated
    # if we discover we need it.  For now, simplicity wins.
    args = config.get('script-args', '').format(**config)
    buildbot_pkg = config.get('buildbot-pkg')
    extra_repo = config.get('extra-repository')
    buildbot_dir = config.get('installdir')

    if extra_repo:
        install_extra_repository(extra_repo)

    if buildbot_pkg:
        log('Installing ' + buildbot_pkg)
        apt_get_install(buildbot_pkg)
        log('Creating initial buildbot slave in ' + buildbot_dir)
        create_slave('temporary', 'temporary', buildbot_dir=buildbot_dir)

    # Some versions of LXC require the user to have a group.  Bug #942850.
    run('addgroup', 'buildbot')
    run('usermod', '--gid', 'buildbot', 'buildbot')

    config_json.set(config)

    retrieve = METHODS.get(method)
    if retrieve and url and path:
        # Make buildbot user have a shell by editing /etc/passwd.
        # Otherwise you cannot ssh as this user, which some scripts
        # need (e.g. those that create lxc containers).  We choose sh as
        # a standard and basic "system" shell.
        run('usermod', '-s', '/bin/sh', 'buildbot')
        sys.exit(handle_script(retrieve, url, path, args))


if __name__ == '__main__':
    log_entry()
    try:
        main()
    finally:
        log_exit()