image-customize

#!/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/>.

import argparse
import os
import re
import sys
import subprocess

from machine.machine_core.constants import BOTS_DIR, TEST_DIR
os.environ["PATH"] = "{0}:{1}".format(os.environ.get("PATH"), BOTS_DIR)

from machine import testvm

parser = argparse.ArgumentParser(
    description='Run command inside or install packages into a Cockpit virtual machine',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-v', '--verbose', action='store_true',
                    help='Display verbose progress details')
parser.add_argument('-i', '--install', action='append', dest="packagelist", default=[],
                    help='Install packages')
parser.add_argument('-I', '--install-command',
                    help="Command used to install packages in machine (default: auto-detect dnf/yum/apt)")
parser.add_argument('-r', '--run-command', action='append', dest="commandlist", default=[],
                    help='Run command inside virtual machine')
parser.add_argument('-s', '--script', action='append', dest="scriptlist", default=[],
                    help='Run selected script inside virtual machine')
parser.add_argument('-u', '--upload', action='append', dest="uploadlist", default=[],
                    help='Upload file/dir to destination file/dir separated by ":" example: -u file.txt:/var/lib')
parser.add_argument('--base-image', help='Base image name, if "image" does not match a standard Cockpit VM image name')
parser.add_argument('--resize', help="Resize the image. Size in bytes with using K, M, or G suffix.")
parser.add_argument('image', help='The image to use (destination name when using --base-image)')
args = parser.parse_args()

if not args.base_image:
    args.base_image = os.path.basename(args.image)

args.base_image = testvm.get_test_image(args.base_image)

# Create the necessary layered image for the build/install


def prepare_install_image(base_image, install_image):
    if "/" not in base_image:
        base_image = os.path.join(testvm.IMAGES_DIR, base_image)
    if "/" not in install_image:
        install_image = os.path.join(os.path.join(TEST_DIR, "images"), os.path.basename(install_image))

    # In vm-customize we don't force recreate images
    if not os.path.exists(install_image):
        install_image_dir = os.path.dirname(install_image)
        os.makedirs(install_image_dir, exist_ok=True)
        base_image = os.path.realpath(base_image)
        qcow2_image = "{0}.qcow2".format(install_image)
        subprocess.check_call(["qemu-img", "create", "-q", "-f", "qcow2",
                               "-o", "backing_file={0},backing_fmt=qcow2".format(base_image), qcow2_image])
        if os.path.lexists(install_image):
            os.unlink(install_image)
        os.symlink(os.path.basename(qcow2_image), install_image)

    if args.resize:
        subprocess.check_call(["qemu-img", "resize", install_image, args.resize])

    return install_image


def run_command(machine_instance, commandlist):
    """Run command inside image"""
    for foo in commandlist:
        try:
            machine_instance.execute(foo, timeout=1800)
        except subprocess.CalledProcessError as e:
            sys.stderr.write("%s\n" % e)
            sys.exit(e.returncode)


def run_script(machine_instance, scriptlist):
    """Run script inside image"""
    for foo in scriptlist:
        if os.path.isfile(foo):
            pname = os.path.basename(foo)
            uploadpath = "/var/tmp/" + pname
            machine_instance.upload([os.path.abspath(foo)], uploadpath)
            machine_instance.execute("chmod a+x %s" % uploadpath)
            try:
                machine_instance.execute(uploadpath, timeout=1800)
            except subprocess.CalledProcessError as e:
                sys.stderr.write("%s\n" % e)
                sys.exit(e.returncode)
        else:
            sys.stderr.write("Bad path to script: %s\n" % foo)


def upload_files(machine_instance, uploadfiles):
    """Upload files/directories inside image"""
    for foo in uploadfiles:
        srcfile, dest = foo.split(":")
        src_absolute = os.path.join(os.getcwd(), srcfile)
        machine_instance.upload([src_absolute], dest)


def install_packages(machine_instance, packagelist, install_command):
    """Install packages into a test image
    It could be done via local rpms or normal package installation
    """

    if not install_command:
        # this will fail if neither is available -- exception is clear enough, this is a developer tool
        out = machine_instance.execute("which dnf || which yum || which apt-get")
        if 'dnf' in out:
            install_command = "dnf install -y"
        elif 'yum' in out:
            install_command = "yum --setopt=skip_missing_names_on_install=False -y install"
        else:
            install_command = "apt-get install -y"

    allpackages = []
    for foo in packagelist:
        if os.path.isfile(foo):
            pname = os.path.basename(foo)
            machine_instance.upload([foo], "/var/tmp/" + pname)
            allpackages.append("/var/tmp/" + pname)
        elif not re.search("/", foo):
            allpackages.append(foo)
        else:
            sys.stderr.write("Bad package name or path: %s\n" % foo)
    if allpackages:
        machine_instance.execute(install_command + " " + ' '.join(allpackages), timeout=1800)


if args.commandlist or args.packagelist or args.scriptlist or args.uploadlist or args.resize:
    if '/' not in args.base_image:
        subprocess.check_call(["image-download", args.base_image])
    machine = testvm.VirtMachine(maintain=True,
                                 verbose=args.verbose, image=prepare_install_image(args.base_image, args.image))
    machine.start()
    machine.wait_boot()
    try:
        if args.uploadlist:
            upload_files(machine, args.uploadlist)
        if args.commandlist:
            run_command(machine, args.commandlist)
        if args.packagelist:
            install_packages(machine, args.packagelist, args.install_command)
        if args.scriptlist:
            run_script(machine, args.scriptlist)
    finally:
        machine.stop()