image-create

#!/usr/bin/env python3
# This file is part of Cockpit.
#
# Copyright (C) 2015 Red Hat, Inc.
#
# Cockpit is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 2.1 of the License, or
# (at your option) any later version.
#
# Cockpit 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Cockpit; If not, see <http://www.gnu.org/licenses/>.

# image-create  -- Make a root image suitable for use with vm-run.
#
# Installs the OS indicated by TEST_OS into the image
# for test machine and tweaks it to be useable with
# vm-run and testlib.py.

import argparse
import os
import shutil
import subprocess
import sys
import time
import tempfile

from machine.machine_core.directories import BOTS_DIR

from machine import testvm

parser = argparse.ArgumentParser(description='Create a virtual machine image')
parser.add_argument('-v', '--verbose', action='store_true', help='Display verbose progress details')
parser.add_argument('-s', '--sit', action='store_true', help='Sit and wait if setup script fails')
parser.add_argument('-n', '--no-save', action='store_true', help='Don\'t save the new image')
parser.add_argument('-u', '--upload', action='store_true', help='Upload the image after creation')
parser.add_argument('--no-build', action='store_true', dest='no_build',
                    help='Don''t build packages and create the vm without build capabilities')
parser.add_argument("--store", default=None, help="Where to send images")
parser.add_argument('image', help='The image to create')
args = parser.parse_args()

# default to --no-build for some images
if args.image in ["continuous-atomic", "rhel-atomic", "fedora-coreos", "services", "selenium", "openshift"]:
    if not args.no_build:
        if args.verbose:
            print("Creating machine without build capabilities based on the image type")
        args.no_build = True


class MachineBuilder:
    def __init__(self, machine):
        tempdir = testvm.get_temp_dir()
        self.machine = machine

        os.makedirs(tempdir, 0o750, exist_ok=True)

        # Use a tmp filename
        self.target_file = self.machine.image_file
        fp, self.machine.image_file = tempfile.mkstemp(dir=tempdir, prefix=self.machine.image, suffix=".qcow2")
        os.close(fp)

    def bootstrap_system(self):
        assert not self.machine._domain

        os.makedirs(self.machine.run_dir, 0o750, exist_ok=True)

        bootstrap_script = os.path.join(testvm.SCRIPTS_DIR, "%s.bootstrap" % (self.machine.image, ))

        if os.path.isfile(bootstrap_script):
            subprocess.check_call([bootstrap_script, self.machine.image_file])
        else:
            raise testvm.Failure("Unsupported OS %s: %s not found." % (self.machine.image, bootstrap_script))

    def run_setup_script(self, script):
        """Prepare a test image further by running some commands in it."""
        self.machine.start()
        try:
            self.machine.wait_boot(timeout_sec=120)
            self.machine.upload([os.path.join(testvm.SCRIPTS_DIR, "lib")], "/var/lib/testvm")
            self.machine.upload([script], "/var/tmp/SETUP")
            self.machine.upload([os.path.join(testvm.SCRIPTS_DIR, "lib", "base")],
                                "/var/tmp/cockpit-base")

            env = {
                "TEST_OS": self.machine.image,
                "DO_BUILD": "0" if args.no_build else "1",
            }
            self.machine.message("run setup script on guest")

            try:
                self.machine.execute(script="/var/tmp/SETUP " + self.machine.image,
                                     environment=env, quiet=not self.machine.verbose, timeout=7200)
                self.machine.execute(command="rm -f /var/tmp/SETUP")

                if self.machine.image == 'openshift':
                    # update our local openshift kube config file to match the new image
                    self.machine.download("/root/.kube/config",
                                          os.path.join(BOTS_DIR, "images/files/openshift.kubeconfig"))

            except subprocess.CalledProcessError as ex:
                if args.sit:
                    sys.stderr.write(self.machine.diagnose())
                    input("Press RET to continue... ")
                raise testvm.Failure("setup failed with code {0}\n".format(ex.returncode))

        finally:
            self.machine.stop(timeout_sec=500)

    def boot_system(self):
        """Start the system to make sure it can boot, then shutdown cleanly
        This also takes care of any selinux relabeling setup triggered
        Don't wait for an ip address during start, since the system might reboot"""
        self.machine.start()
        try:
            self.machine.wait_boot(timeout_sec=120)
        finally:
            self.machine.stop(timeout_sec=120)

    def build(self):
        self.bootstrap_system()

        # gather the scripts, separated by reboots
        script = os.path.join(testvm.SCRIPTS_DIR, "%s.setup" % (self.machine.image, ))

        if not os.path.exists(script):
            return

        self.machine.message("Running setup script %s" % (script))
        self.run_setup_script(script)

        tries_left = 3
        successfully_booted = False
        while tries_left > 0:
            try:
                # make sure we can boot the system
                self.boot_system()
                successfully_booted = True
                break
            except:
                # we might need to wait for the image to become available again
                # accessing it in maintain=True mode successively can trigger qemu errors
                time.sleep(3)
                tries_left -= 1
        if not successfully_booted:
            raise testvm.Failure("Unable to verify that machine boot works.")

    def save(self):
        data_dir = testvm.get_images_data_dir()

        os.makedirs(data_dir, 0o750, exist_ok=True)

        if not os.path.exists(self.machine.image_file):
            raise testvm.Failure("Nothing to save.")

        partial = os.path.join(data_dir, self.machine.image + ".partial")

        # Copy image via convert, to make it sparse again
        subprocess.check_call(["qemu-img", "convert", "-c", "-O", "qcow2", self.machine.image_file, partial])

        # Hash the image here
        (sha, x1, x2) = subprocess.check_output(["sha256sum", partial], universal_newlines=True).partition(" ")
        if not sha:
            raise testvm.Failure("sha256sum returned invalid output")

        name = self.machine.image + "-" + sha + ".qcow2"
        data_file = os.path.join(data_dir, name)
        shutil.move(partial, data_file)

        # Remove temp image file
        os.unlink(self.machine.image_file)

        # Update the images symlink
        if os.path.islink(self.target_file):
            os.unlink(self.target_file)
        os.symlink(name, self.target_file)

        # Handle alternate images data directory
        image_file = os.path.join(testvm.IMAGES_DIR, name)
        if not os.path.exists(image_file):
            os.symlink(os.path.abspath(data_file), image_file)


try:
    testvm.VirtMachine.memory_mb = 2048
    machine = testvm.VirtMachine(verbose=args.verbose, image=args.image, maintain=True)
    builder = MachineBuilder(machine)
    builder.build()
    if not args.no_save:
        print("Saving...")
        builder.save()
        if args.upload:
            print("Uploading...")
            cmd = [os.path.join(BOTS_DIR, "image-upload")]
            if args.store:
                cmd += ["--store", args.store]
            cmd += [args.image]
            subprocess.check_call(cmd)

except testvm.Failure as ex:
    sys.stderr.write("image-create: %s\n" % ex)
    sys.exit(1)