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 tempfile

from lib.constants import BOTS_DIR, DEFAULT_IDENTITY_PUB_FILE
from machine import testvm
from machine.machine_core import ssh_connection

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('--capture-console', action='store_true', help='Capture VM console to stderr')
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 ["fedora-coreos", "services"]:
    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):
        self.machine = machine
        self.iso = self.machine.image.endswith("boot")
        self.payload = self.machine.image.endswith("payload")
        if self.iso:
            self.suffix = ".iso"
        elif self.payload:
            self.suffix = ".tar.gz"
        else:
            self.suffix = ".qcow2"

        # Use /var/tmp/ as this is going to be a huge file; /tmp/ is commonly tmpfs
        self.target_file = self.machine.image_file
        fp, self.machine.image_file = tempfile.mkstemp(dir="/var/tmp/", prefix=self.machine.image, suffix=self.suffix)
        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()

        # Edge management does not allow setting key for root account but rest of this script depends on it
        if self.machine.image == "rhel4edge":
            admin_s = ssh_connection.SSHConnection(user="admin",
                                                   address=self.machine.ssh_address,
                                                   ssh_port=self.machine.ssh_port,
                                                   identity_file=self.machine.identity_file)
            admin_s.wait_boot(timeout_sec=120)
            with open(DEFAULT_IDENTITY_PUB_FILE) as pub:
                admin_s.execute(f"echo '{pub.read().strip()}' | sudo tee -a /root/.ssh/authorized_keys")

        try:
            self.machine.wait_boot(timeout_sec=300)
            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",
                "SERVER_REPO_URL": os.environ.get("SERVER_REPO_URL", ""),
                "EXTRAS_REPO_URL": os.environ.get("EXTRAS_REPO_URL", ""),
                "BASEOS_REPO_URL": os.environ.get("BASEOS_REPO_URL", ""),
                "APPSTREAM_REPO_URL": os.environ.get("APPSTREAM_REPO_URL", ""),
            }
            self.machine.message("run setup script on guest")

            try:
                self.machine.execute(f"/var/tmp/SETUP {self.machine.image}",
                                     environment=env, stdout=None, quiet=not self.machine.verbose, timeout=7200)
                self.machine.execute("rm -f /var/tmp/SETUP")
            except subprocess.CalledProcessError as ex:
                if args.sit:
                    sys.stderr.write(self.machine.diagnose())
                    input("Press RET to continue... ")
                raise testvm.Failure(f"setup failed with code {ex.returncode}\n")

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

    def boot_system(self):
        """Start the system to make sure it can boot, then shutdown cleanly"""

        self.machine.start()
        # avoid too long boot times -- they cause painfully long tests, and are a bug
        try:
            timeout = 30
            if self.machine.image == "fedora-rawhide":
                # except with fedora-rawhide, which is slow by design:
                # https://lists.fedoraproject.org/archives/list/devel@lists.fedoraproject.org/thread/HRJT4UF35MRJ7C4ZXVHBFXB4BTW64YSO/
                timeout = 120
            self.machine.wait_boot(timeout_sec=timeout)
        finally:
            self.machine.stop(timeout_sec=30)

    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)

        # make sure we can boot the system
        self.boot_system()

    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.")

        if not self.iso and not self.payload:
            rebuild = os.path.join(data_dir, self.machine.image + ".rebuild")

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

        # Hash the image here
        (sha, x1, x2) = subprocess.check_output(
            ["sha256sum", self.machine.image_file if (self.iso or self.payload) else rebuild], universal_newlines=True
        ).partition(" ")
        if not sha:
            raise testvm.Failure("sha256sum returned invalid output")

        name = self.machine.image + "-" + sha + self.suffix
        data_file = os.path.join(data_dir, name)
        shutil.move(self.machine.image_file if (self.iso or self.payload) else rebuild, data_file)

        if not self.iso and not self.payload:
            # 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:
    if args.image == 'services':
        # deploying candlepin needs oodles of memory
        memory_mb = 3072
    else:
        memory_mb = 2048

    machine = testvm.VirtMachine(verbose=args.verbose,
                                 capture_console=args.capture_console,
                                 image=args.image,
                                 memory_mb=memory_mb,
                                 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"), '--prune-s3']
            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)