~joetalbott/snappy-proposed-image-builder/add_versions

# snappy-proposed-image-builder
# Copyright (C) 2015 Canonical
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""Business logic or the service lives here."""

import contextlib
import glob
import logging
import os
import subprocess
import tempfile

from uservice_utils.queue import MessageActions

from snappy_proposed_image_builder.constants import LOGGING_EXTRA
from snappy_proposed_image_builder.cloud import get_glance_client
from snappy_proposed_image_builder.utils import (
    check_call,
    BetterCalledProcessError,
)


logger = logging.getLogger(__name__)


class RootFSError(Exception):
    pass


class ImageBuilderWorker(object):

    """A worker callable that contains all our main logic."""

    def __init__(self, config, builder):
        self._config = config
        self.publish_results = builder

    def __call__(self, payload):
        extra = LOGGING_EXTRA.copy()
        extra.update(payload)
        try:
            release = payload['release']
            channel = payload['channel']
            device = payload['device']
            image_binaries = payload['image_binaries']
            source_version = payload['source_version']
        except KeyError as e:
            logger.error(
                "Unable to deserialize message payload - "
                "rejecting message: %s",
                e,
                extra=extra
            )
            return MessageActions.Retry

        (image_version, rootfs_version, device_version) = _get_image_versions(
            release, channel, device)

        payload['image_version'] = image_version
        payload['rootfs_version'] = rootfs_version
        payload['device_version'] = device_version

        with tempfile.TemporaryDirectory() as tmpdir:
            logger.info("Beginning rootfs download.", extra=extra)
            try:
                rootfs_path = download_rootfs(
                    image_version,
                    release,
                    channel,
                    device,
                    tmpdir
                )
            except subprocess.CalledProcessError as e:
                logger.error(
                    "Unable to download core image: %s",
                    e,
                    extra=extra
                )
                return MessageActions.Retry

            logger.info("Adding packages from proposed")
            MAX_RETRIES = 3
            count = MAX_RETRIES
            while count > 0:
                new_rootfs_path = add_proposed_packages(rootfs_path,
                                                        image_binaries,
                                                        source_version)
                if new_rootfs_path is not None:
                    break

                # try removing the rootfs as it sometimes is corrupted.
                # XXX: find out why it's getting corrupted.
                logger.info("Retrying rootfs download.", extra=extra)
                os.remove(rootfs_path)
                try:
                    rootfs_path = download_rootfs(
                        image_version,
                        channel,
                        device,
                        tmpdir
                    )
                except subprocess.CalledProcessError as e:
                    logger.error(
                        "Unable to download core image: %s",
                        e,
                        extra=extra
                    )
                    return MessageActions.Retry

                count -= 1
            if count == 0:
                logger.error(
                    "Unable to add packages after %d tries: %s",
                    MAX_RETRIES,
                    image_binaries,
                    source_version,
                )
                return MessageActions.Retry

            try:
                logger.info("Building the image")
                image_path = build_image(
                    image_version,
                    new_rootfs_path,
                    release,
                    channel,
                    device,
                    tmpdir
                )
            except subprocess.CalledProcessError as e:
                logger.error(
                    "Unable to build the image: %s",
                    e,
                    extra=extra
                )
                return MessageActions.Retry
            logger.info("Ubuntu Core image build OK.", extra=extra)

            try:
                nova_image_path = convert_nova_image(image_path)
            except subprocess.CalledProcessError as e:
                logger.error(
                    "Unable to convert core image to qcow2 image: %s",
                    e,
                    extra=extra)
                return MessageActions.Retry
            logger.info("Image converted to qcow2 OK.", extra=extra)

            try:
                logger.info("Beginning image upload to glance...",
                            extra=extra)
                glance_image_id = upload_image_to_glance(
                    nova_image_path,
                    self._config
                )
            except Exception as e:
                logger.error(
                    "Unable to upload image to glance: %s",
                    e,
                    extra=extra
                )
                return MessageActions.Retry
            logger.info("Image uploaded to glance OK.", extra=extra)

        payload['nova_image_id'] = glance_image_id
        self.publish_results(payload)
        logger.info("Processing completed.", extra=extra)
        return MessageActions.Acknowledge


def download_rootfs(name, release, channel, device, tmpdir):
    """Download the ubuntu code image, return a path to it on disk."""
    image_path = os.path.join(tmpdir, 'snappy-proposed-{}.img'.format(name))

    cmd = ['sudo',
           './ubuntu-device-flash',
           '--revision', name,
           '--download-only',
           'core',
           '--device', device,
           '--channel', channel,
           '--size', '3',
           '-o', image_path,
           '--developer-mode',
           '--cloud',
           release]
    check_call(cmd)

    cmd = ['sudo',
           'rm',
           '-rf',
           os.path.join(tmpdir, '.gnupg')]
    check_call(cmd)
    output = glob.glob('ubuntu-*.tar.xz')
    return output[0]


@contextlib.contextmanager
def mount_proc(chroot_base):
    check_call(['sudo', 'mount', '-t', 'proc', 'none',
                os.path.join(chroot_base, 'proc')])
    try:
        yield
    finally:
        check_call(['sudo', 'umount', os.path.join(chroot_base, 'proc')])


@contextlib.contextmanager
def mount_dev_pts(chroot_base):
    check_call(['sudo', 'mount', '-t', 'devpts', '-o', 'gid=5,mode=620',
                'none', os.path.join(chroot_base, 'dev', 'pts')])
    try:
        yield
    finally:
        check_call(['sudo', 'umount', os.path.join(chroot_base, 'dev', 'pts')])


@contextlib.contextmanager
def mount_sys(chroot_base):
    check_call(['sudo', 'mount', '-t', 'sysfs', 'none',
                os.path.join(chroot_base, 'sys')])
    try:
        yield
    finally:
        check_call(['sudo', 'umount', os.path.join(chroot_base, 'sys')])


@contextlib.contextmanager
def mount_dev_shm(chroot_base):
    check_call(['sudo', 'mount', '-t', 'tmpfs', 'none',
                os.path.join(chroot_base, 'dev', 'shm')])
    try:
        yield
    finally:
        check_call(['sudo', 'umount', os.path.join(chroot_base, 'dev', 'shm')])


def add_proposed_packages(rootfs_path, image_binaries, source_version):

    new_rootfs_path = rootfs_path.replace('.tar.xz', '-modified.tar.xz')

    ROOTFS_BASE_DIR = 'system'

    # untar the rootfs
    cmd = ['sudo', 'tar', 'xJf', rootfs_path]
    try:
        check_call(cmd)
    except BetterCalledProcessError as e:
        logger.error(e)
        return None

    # sanity check that the base directory exists
    if not os.path.exists('system'):
        logger.error("Failed to untar the rootfs: %s", rootfs_path)
        return None

    chroot_base = os.path.join(os.getcwd(), ROOTFS_BASE_DIR)

    # chroot into the rootfs and install packages
    cmd = ['sudo',  'chroot', chroot_base,
           'lsb_release', '-cs']
    output = subprocess.check_output(cmd)
    series = output.decode("utf-8").strip()
    proposed_pocket = '{}-proposed'.format(series)
    deb_url = "http://ftpmaster.internal/ubuntu/"
    proposed_line = ("deb {} {} restricted main universe".format(
        deb_url, proposed_pocket))
    proposed_line_2 = ""

    apt_sources_path = os.path.join(chroot_base, 'etc', 'apt', 'sources.list')
    with open(apt_sources_path, 'r') as fp:
        for line in fp.readlines():
            if line.startswith('deb '):
                data = line.split()
                if len(data) > 1:  # assumes a properly formatted source entry
                    proposed_line_2 = proposed_line.replace(deb_url, data[1])
                    break

    new_apt_sources_file = os.path.join(chroot_base, 'etc', 'apt',
                                        'sources.list.d',
                                        'add_proposed_pocket')
    with tempfile.NamedTemporaryFile(mode='wt') as fp:
        fp.write(proposed_line)
        fp.write('\n')
        fp.write(proposed_line_2)
        fp.flush()

        cmd = ['sudo', 'cp', fp.name, new_apt_sources_file]
        check_call(cmd)

        cmd = ['sudo', 'chmod', '+r',  new_apt_sources_file]
        check_call(cmd)

    with contextlib.ExitStack() as mount_stack:
        mount_stack.enter_context(mount_proc(chroot_base))
        mount_stack.enter_context(mount_sys(chroot_base))
        mount_stack.enter_context(mount_dev_pts(chroot_base))
        mount_stack.enter_context(mount_dev_shm(chroot_base))

        check_call(['sudo', 'cp', '/etc/hosts', '/etc/hostname',
                    '/etc/resolv.conf', os.path.join(chroot_base, 'etc')])

        cmd = ['sudo',  'chroot', chroot_base,
               '/usr/bin/apt-get', 'update']
        check_call(cmd)

        for package_name in image_binaries:
            pkg_version_cmd = ['sudo', 'chroot', chroot_base, 'dpkg-query',
                               '--show', package_name]

            old_version = subprocess.check_output(pkg_version_cmd)
            old_version = old_version.strip().decode('utf-8')

            if old_version == "":
                raise RootFSError(
                    "package not installed on original image: {}".format(
                        package_name))

            cmd = ['sudo',  'chroot', chroot_base,
                   '/usr/bin/apt-get', 'install', '--yes',
                   "{}={}".format(package_name, source_version)]
            try:
                check_call(cmd)
            except subprocess.CalledProcessError as e:
                raise RootFSError(
                    "Failed to install package: {}\n: {}".format(
                        package_name, e))

            new_version = subprocess.check_output(pkg_version_cmd)

            new_version = new_version.strip().decode('utf-8')
            expected_version = "{}\t{}".format(package_name,
                                               source_version)

            if new_version != expected_version:
                raise RootFSError(
                    "new version is not the same as "
                    "the expected version, {} == {}".format(
                        new_version, expected_version))

    # tar the rootfs back up.
    cmd = ['sudo', 'tar', 'cJf', new_rootfs_path, ROOTFS_BASE_DIR]
    check_call(cmd)

    return new_rootfs_path


def build_image(name, rootfs_path, release, channel, device, tmpdir):
    """Download the ubuntu code image, return a path to it on disk."""
    image_path = os.path.join(tmpdir, 'snappy-proposed-{}.img'.format(name))
    cmd = ['sudo',
           './ubuntu-device-flash',
           '--revision', name,
           'core',
           '--device', device,
           '--channel', channel,
           '--size', '3',
           '--image-part', rootfs_path,
           '-o', image_path,
           '--developer-mode',
           '--cloud',
           release]
    check_call(cmd)

    cmd = ['sudo',
           'rm',
           '-rf',
           os.path.join(tmpdir, '.gnupg')]
    check_call(cmd)

    return image_path


def convert_nova_image(image_path):
    """
    Convert a core image to a nova image, return path to the converted image.
    """
    converted_image_path = '{}-cloud.img'.format(image_path.strip('.img'))
    cmd = ['qemu-img',
           'convert',
           '-f', 'raw',
           '-O', 'qcow2',
           image_path,
           converted_image_path]
    check_call(cmd)

    return converted_image_path


def upload_image_to_glance(nova_image_path, config):
    """
    Upload the nova image to glance, returning the id of the image in glance.
    """
    extra = LOGGING_EXTRA.copy()
    glance = get_glance_client(config)
    image_name = os.path.basename(nova_image_path)
    image = glance.images.create(
        name=image_name,
        data=open(nova_image_path, 'rb'),
        disk_format='qcow2',
        container_format='bare',
    )
    logger.info("Uploaded %s with ID %s.", image_name, image.id, extra=extra)
    return image.id


def _get_version_string_output(release, channel, device):
    """Obtains a bytestring of the images info from the core image server"""
    extra = LOGGING_EXTRA.copy()
    cmd = [
        './ubuntu-device-flash',
        'query',
        '--show-image',
        '--channel=ubuntu-core/{}/{}'.format(release, channel),
        '--device={}'.format(device),
    ]
    latest_image_info = b''
    try:
        latest_image_info = subprocess.check_output(cmd)
    except subprocess.CalledProcessError as e:
        logger.error(e, exc_info=True, extra=extra)
    finally:
        return latest_image_info


def _get_image_versions(release, channel, device):
    latest_image_info = _get_version_string_output(release, channel, device)
    image_version = ''
    rootfs_version = ''
    dev_version = ''
    for info in latest_image_info.decode('utf-8').split("\n"):
        if 'Description' in info:
            for l in info.split(','):
                if 'version' in l:
                    image_version = l.split('version=')[1]
                    continue
                if 'ubuntu' in l:
                    rootfs_version = l.split('ubuntu=')[1]
                    continue
                if 'raw-device' in l:
                    dev_version = l.split('raw-device=')[1]
                    continue

    return (image_version, rootfs_version, dev_version)