~joetalbott/uci-engine/user_auth

#!/usr/bin/env python
# Ubuntu CI Engine
# Copyright 2014 Canonical Ltd.

# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License version 3, as
# published by the Free Software Foundation.

# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranties of
# MERCHANTABILITY, SATISFACTORY QUALITY, or FITNESS FOR A PARTICULAR
# PURPOSE.  See the GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import re
import os
import sys
import subprocess
from keystoneclient.v2_0 import client as ksclient
import glanceclient
import glanceclient.common.utils
import glanceclient.common.progressbar
import swiftclient
from operator import itemgetter
import tempfile
import shutil
import argparse
import contextlib
import distutils.spawn
import time
import textwrap


def _strip_version(endpoint):
    """Strip version from the last component of endpoint if present."""

    # Get rid of trailing '/' if present
    if endpoint.endswith('/'):
        endpoint = endpoint[:-1]
    url_bits = endpoint.split('/')
    # regex to match 'v1' or 'v2.0' etc
    if re.match('v\d+\.?\d*', url_bits[-1]):
        endpoint = '/'.join(url_bits[:-1])
    return endpoint


def _get_ksclient(**kwargs):
    """Get an endpoint and auth token from Keystone.

    :param username: name of user
    :param password: user's password
    :param tenant_id: unique identifier of tenant
    :param tenant_name: name of tenant
    :param auth_url: endpoint to authenticate against
    """
    return ksclient.Client(username=kwargs.get('username'),
                           password=kwargs.get('password'),
                           tenant_id=kwargs.get('tenant_id'),
                           tenant_name=kwargs.get('tenant_name'),
                           auth_url=kwargs.get('auth_url'),
                           cacert=kwargs.get('cacert'),
                           insecure=kwargs.get('insecure'))


def _get_endpoint(client, **kwargs):
    """Get an endpoint using the provided keystone client."""
    endpoint_kwargs = {
        'service_type': kwargs.get('service_type') or 'image',
        'endpoint_type': kwargs.get('endpoint_type') or 'publicURL',
    }

    if kwargs.get('region_name'):
        endpoint_kwargs['attr'] = 'region'
        endpoint_kwargs['filter_value'] = kwargs.get('region_name')

    endpoint = client.service_catalog.url_for(**endpoint_kwargs)
    return _strip_version(endpoint)


def _get_glanceclient():
    kw = {
        'username': os.environ['OS_USERNAME'],
        'password': os.environ['OS_PASSWORD'],
        'tenant_name': os.environ['OS_TENANT_NAME'],
        'auth_url': os.environ['OS_AUTH_URL'],
        'cacert': os.environ.get('OS_CACERT'),
        'insecure': False,
        'region_name': os.environ['OS_REGION_NAME'],
    }
    _ksclient = _get_ksclient(**kw)
    token = _ksclient.auth_token
    endpoint = _get_endpoint(_ksclient, **kw)
    kw = {
        'token': token,
        'insecure': False,
        'cacert': os.environ.get('OS_CACERT'),
        'ssl_compression': True,
    }
    client = glanceclient.Client('1', endpoint, **kw)
    return client


def _get_swiftclient():
    auth_ver = '/v1.' in os.environ['OS_AUTH_URL'] and '1.0' or '2.0'
    kw = {
        'user': os.environ['OS_USERNAME'],
        'key': os.environ['OS_PASSWORD'],
        'authurl': os.environ['OS_AUTH_URL'],
        'auth_version': auth_ver,
        'os_options': {
            'tenant_name': os.environ['OS_TENANT_NAME'],
            'region_name': os.environ['OS_REGION_NAME'],
            'insecure': False,
        }
    }
    return swiftclient.Connection(**kw)


def _precise_images(client):
    '''Returns a generator of tuples of the form:
       (the image identifier, the image name)'''

    search_prefix = 'ubuntu-released/ubuntu-precise-12.04-amd64'
    for image in client.images.list():
        if image.name.startswith(search_prefix):
            yield (image.id, image.name)


def _most_recent_precise_image(gen):
    '''Returns the image identifier from a generator of tuples of the form:
       (the image identifier, the image name)'''

    return max(gen, key=itemgetter(1))[0]


def _nbd_present():
    '''Return True if /dev/nbd0 exists.'''
    # We could check for nbd in /proc/modules, but maybe it's built-in.
    return os.path.exists('/dev/nbd0')


def _chroot(mountpoint, cmd):
    return subprocess.check_output(['chroot', mountpoint] + cmd)


def _download_most_recent_precise_image(client):
    most_recent = _most_recent_precise_image(_precise_images(client))
    with open(most_recent, 'w') as fp:
        data = client.images.data(most_recent, do_checksum=True)
        for datum in data:
            fp.write(datum)
    return most_recent


def add_ci_ppa(mountpoint):
    '''Add the CI engineering PPA to the image.'''

    pkg = 'software-properties-common'
    _apt_get_install(mountpoint, pkgs=[pkg])
    ppa = 'ppa:canonical-ci-engineering/ci-airline-phase-0'
    _chroot(mountpoint, ['apt-add-repository', '-y', ppa])


def add_cloud_archive(mountpoint):
    '''Add the Ubuntu cloud archive to the image, for in-cloud apt installs.'''

    sed = 's,archive.ubuntu.com,az3.clouds.archive.ubuntu.com,'
    sources_list = '/etc/apt/sources.list'
    _chroot(mountpoint, ['sed', '-i', sed, sources_list])


def add_cloud_tools(mountpoint):
    '''Add the cloud tools archive to the image, for mongodb-server.'''

    tools = ('deb http://ubuntu-cloud.archive.canonical.com/ubuntu '
             'precise-updates/cloud-tools main')
    _chroot(mountpoint, ['apt-add-repository', '-y', tools])
    cmd = ['apt-key', 'adv', '--keyserver', 'keyserver.ubuntu.com',
           '--recv-keys', '5EDB1B62EC4926EA']
    _chroot(mountpoint, cmd)


@contextlib.contextmanager
def no_service_start(mountpoint):
    policyfile = textwrap.dedent("""\
        #!/bin/sh
        while true; do
        case "$1" in
            -*) shift ;;
            makedev) exit 0;;
            x11-common) exit 0;;
            *) exit 101;;
        esac
        done""")
    with open('%s/usr/sbin/policy-rc.d' % mountpoint, 'w') as f:
        f.write(policyfile)
    os.chmod('%s/usr/sbin/policy-rc.d' % mountpoint, 0o755)
    try:
        yield {}
    finally:
        os.unlink('%s/usr/sbin/policy-rc.d' % mountpoint)


@contextlib.contextmanager
def setup_networking(mountpoint):
    '''Copy in the host system's resolv.conf, so network requests are
    routable.'''

    resolv = '/etc/resolv.conf'
    chroot_resolv = os.path.join(mountpoint, 'etc/resolv.conf')
    shutil.move(chroot_resolv, '%s.old' % chroot_resolv)
    shutil.copy(resolv, chroot_resolv)
    try:
        yield {}
    finally:
        shutil.move('%s.old' % chroot_resolv, chroot_resolv)


@contextlib.contextmanager
def setup_bindmounts(mountpoint):
    '''Ensure the chroot has access to /dev, /proc, and /sys for apt.'''

    mounts = ('/proc', '/sys', '/dev', '/dev/pts')
    try:
        for mount in mounts:
            bmount = mountpoint + mount
            subprocess.check_call(['mount', '--bind', mount, bmount])
        yield {}
    finally:
        for mount in reversed(mounts):
            bmount = mountpoint + mount
            subprocess.check_call(['umount', bmount])


@contextlib.contextmanager
def mounted_image(img):
    # Ensure it's not connected.
    subprocess.check_call(['qemu-nbd', '-d', '/dev/nbd0'])
    subprocess.check_call(['qemu-nbd', '-c', '/dev/nbd0', img])
    mountpoint = tempfile.mkdtemp()
    subprocess.check_call(['mount', '/dev/nbd0p1', mountpoint])
    try:
        yield mountpoint
    finally:
        subprocess.check_call(['umount', mountpoint])
        os.rmdir(mountpoint)
        subprocess.check_call(['qemu-nbd', '-d', '/dev/nbd0'])


def resize_image(img):
    '''Make the provided image 5GB larger.'''
    subprocess.check_call(['qemu-img', 'resize', img, '+5G'])
    subprocess.check_call(['qemu-nbd', '-d', '/dev/nbd0'])
    subprocess.check_call(['qemu-nbd', '-c', '/dev/nbd0', img])
    try:
        subprocess.check_call(['e2fsck', '-f', '/dev/nbd0p1'])
        # Resize the partition to use the new available space.
        subprocess.check_call(['resize2fs', '/dev/nbd0p1'])
    finally:
        subprocess.check_call(['qemu-nbd', '-d', '/dev/nbd0'])


def install_packages(mountpoint):
    # http://bazaar.launchpad.net/~go-bot/juju-core/trunk/view/head:/environs/cloudinit/cloudinit.go#L210
    juju_pkgs = [
        'git',
        'curl',
        'cpu-checker',
        'bridge-utils',
        'rsyslog-gnutls',
    ]
    pkgs = [
        'bzr',
        'dput',
        'qemu-utils',
        'gunicorn',
        'python-dput',
        'python-jenkins',
        'python-swiftclient',
        'python-novaclient',
        'python-glanceclient',
        'python-amqplib',
        'python-oauth',
        'python-launchpadlib',
        'python-django',
        'python-tastypie',
        'python-django-south',
        'python-yaml',
        'python-lazr.enum',
        'python-requests',
        'python-jinja2',
        'python-txstatsd',
        'python-tz',
        'python-gnupg',
        'python-restish',
        'rabbitmq-server',
        'python-pip',
        'mercurial',
        'git-core',
        'subversion',
        'gettext',
        'postgresql-9.1',
        'postgresql-contrib-9.1',
        'python-psutil',
        'python-psycopg2',
        'pwgen',
        'postgresql-client',
        'python-support',
        'pgtune',
        'postgresql-9.1-debversion',
        'postgresql-plpython-9.1',
        'python-dnspython',
    ]
    pkgs += juju_pkgs
    _apt_get_install(mountpoint, pkgs=pkgs)

    args = ['-t', 'precise-updates/cloud-tools']
    _apt_get_install(mountpoint, args, ['mongodb-server'])

    _chroot(mountpoint, ['apt-get', 'clean'])


def _apt_get_install(mountpoint, args=None, pkgs=None):
    apt_get_install = ['apt-get', '-qy']
    if args:
        apt_get_install += args
    apt_get_install += ['install'] + pkgs
    try:
        _chroot(mountpoint, apt_get_install)
    except subprocess.CalledProcessError:
        for path in ('var/log/dpkg.log', 'var/log/apt/term.log'):
            with open(os.path.join(mountpoint, path)) as fp:
                print '%s:' % path
                print '-' * (len(path) + 1)
                print fp.read()
        print 'Failed to apt-get install some packages.'
        raise SystemExit


def _add_members_to_image(image_id, members):
    client = _get_glanceclient()
    for member_id in members:
        client.image_members.create(image_id, member_id)


def _check_bake_deps():
    '''Ensure we have qemu-nbd installed, from the qemu-utils package.'''
    return distutils.spawn.find_executable('qemu-nbd') is not None


def _create_metadata(image_id, container='images'):
    meta = tempfile.mkdtemp()
    client = _get_swiftclient()
    try:
        cmd = ['juju-metadata', 'generate-image', '-i', image_id, '-d', meta]
        with open('/dev/null', 'w') as devnull:
            subprocess.check_call(cmd, stdout=devnull)
        headers = {'x-container-read': '.r:*'}
        client.put_container(container, headers=headers)
        # Not to be confused with container above, this is a fixed location
        # that juju-metadata places its files in.
        base = len(os.path.join(meta, 'images')) + 1
        for root, dirs, files in os.walk(os.path.join(meta, 'images')):
            for filename in files:
                filename = os.path.join(root, filename)
                obj = filename[base:]
                with open(os.path.join(root, filename), 'rb') as fp:
                    kw = {
                        'container': container,
                        'obj': obj,
                        'contents': fp.read(),
                    }
                    client.put_object(**kw)
    finally:
        shutil.rmtree(meta)
    print 'Add the following parameter to environments.yaml'
    print 'image-metadata-url: %s/%s' % (client.url, container)


def _parse_args():
    desc = 'Bake an image with dependencies.'
    epilog = '''example:
  source ~/.canonistack/novarc
  # Download the latest daily precise amd64 image.
  sudo -E {cmd} get
  IMAGE="$(ls -tr | tail -n1)"

  # Remaster the cloud image with our and Juju's dependencies, add the
  # cloud mirror, and upgrade.
  sudo -E {cmd} bake $IMAGE
  source ~/.hpcloud-rc

  # Upload the image to HP Cloud and share it with the team.
  sudo -E {cmd} put $IMAGE --members 11086019986478 \\
  11206487910601 11269895438533 11289530460295 11293005633044 \\
  11296464907126 11433916157270 11597203075020 11630049285977 \\
  11740916806557 11845411957545 11885335739817 11922491647292 \\
  11928809798385 11935418143898 11859792530542 11904293261511 \\
  11685836270002

  # Update the simplestreams metadata for Juju
  export OS_TENANT_NAME=juju_tools
  sudo -E {cmd} use $IMAGE'''.format(cmd=sys.argv[0])
    kw = {
        'formatter_class': argparse.RawDescriptionHelpFormatter,
        'description': desc,
        'epilog': epilog,
    }
    parser = argparse.ArgumentParser(**kw)
    kw = {'dest': 'operation', 'help': 'sub-command -h'}
    subparsers = parser.add_subparsers(**kw)
    subparsers.add_parser('get', help='get -h')
    bake = subparsers.add_parser('bake', help='bake -h')
    bake.add_argument('image', help='path to the image to remaster')
    put = subparsers.add_parser('put', help='put -h')
    put.add_argument('image', help='path to the image to remaster')
    kw = {
        'help': 'list of tenants to give access to the image',
        'required': True,
        'nargs': '+',
        'metavar': 'tenant_id',
    }
    put.add_argument('--members', **kw)
    use = subparsers.add_parser('use', help='use -h')
    h = 'container to put metadata in (default: images)'
    use.add_argument('--container', default='images', help=h)
    use.add_argument('image', help='path to the image to create metadata for')
    return parser.parse_args()


def get(*args, **kw):
    '''Download the latest precise image to the current working directory.'''

    # TODO we should get this from
    # http://cloud-images.ubuntu.com/releases/precise/release/ubuntu-12.04-server-cloudimg-amd64-disk1.img
    # instead.
    client = _get_glanceclient()
    try:
        print _download_most_recent_precise_image(client)
        return 0
    except glanceclient.exc.CommunicationError as e:
        print str(e)
        print >>sys.stderr, '\nDo you have a ssh tunnel to Glance?'
        return 1


def bake(*args, **kw):
    '''Remaster the image with dependencies.'''

    if not _check_bake_deps():
        print >>sys.stderr, 'Please install qemu-utils first.'
        return 1

    img = tempfile.mktemp()
    # Make updating this image atomic.
    shutil.copy(kw['image'], img)
    try:
        resize_image(img)
        with mounted_image(img) as mountpoint:
            with setup_bindmounts(mountpoint):
                with setup_networking(mountpoint):
                    with no_service_start(mountpoint):
                        add_ci_ppa(mountpoint)
                        add_cloud_archive(mountpoint)
                        add_cloud_tools(mountpoint)
                        _chroot(mountpoint, ['apt-get', 'update'])
                        install_packages(mountpoint)
                        _chroot(mountpoint, ['apt-get', '-y', 'upgrade'])
    except:
        os.unlink(img)
        raise
    shutil.move(img, kw['image'])
    return 0


def put(*args, **kw):
    with mounted_image(kw['image']) as mountpoint:
        release = _chroot(mountpoint, ['lsb_release', '-sd']).strip('\n')
        arch = _chroot(mountpoint, ['dpkg', '--print-architecture'])
        arch = arch.strip('\n')
        p = os.path.join(mountpoint, 'var/log/dpkg.log')
        if os.path.exists(p):
            mtime = os.path.getmtime(p)
        else:
            mtime = os.path.getmtime(os.path.join(mountpoint, 'var/log'))
    image_date = time.strftime('%Y%m%d', time.gmtime(mtime))
    name = '%s (%s %s) - CI Engineering' % (release, arch, image_date)
    client = _get_glanceclient()
    image = client.images.create(name=name)
    kwargs = {
        'data': open(kw['image'], 'rb'),
        'is_public': False,
        'container_format': 'bare',
        'disk_format': 'qcow2',
    }

    filesize = glanceclient.common.utils.get_file_size(kwargs['data'])
    args = (kwargs['data'], filesize)
    kwargs['data'] = glanceclient.common.progressbar.VerboseFileWrapper(*args)
    image.update(**kwargs)
    _add_members_to_image(image.id, kw['members'])
    print image.id
    return 0


def use(*args, **kw):
    _create_metadata(kw['image'], kw['container'])


def main():
    if os.getuid() != 0:
        print >>sys.stderr, 'This application needs to be run under `sudo -E`.'
        return 1

    if not _nbd_present():
        m = 'Please run `sudo modprobe nbd` before running this command.'
        print >>sys.stderr, m
        return 1

    args = _parse_args()
    # Call get(), bake(), or put() with the parsed arguments as keyword
    # arguments to the function.
    return globals()[args.operation](**vars(args))


if __name__ == '__main__':
    sys.exit(main())