lorax/src/sbin/livemedia-creator

#!/usr/bin/python
#
# Live Media Creator
#
# Copyright (C) 2011-2014  Red Hat, Inc.
#
# 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 2 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/>.
#
import logging
log = logging.getLogger("livemedia-creator")
program_log = logging.getLogger("program")
pylorax_log = logging.getLogger("pylorax")

import argparse
import os
import shutil
import sys
import tempfile

# Use pykickstart to calculate disk image size
from pykickstart.parser import KickstartParser
from pykickstart.version import makeVersion, RHEL7

from pylorax import vernum
from pylorax.creator import DRACUT_DEFAULT, mount_boot_part_over_root
from pylorax.creator import make_appliance, make_image, make_livecd, make_live_images
from pylorax.creator import make_runtime, make_squashfs
from pylorax.imgutils import copytree
from pylorax.imgutils import Mount, PartitionMount, umount
from pylorax.installer import InstallError
from pylorax.sysutils import joinpaths

VERSION = "{0}-{1}".format(os.path.basename(sys.argv[0]), vernum)


# no-virt mode doesn't need libvirt, so make it optional
try:
    import libvirt
except ImportError:
    libvirt = None

def lorax_parser():
    """ Return the ArgumentParser for lorax"""

    parser = argparse.ArgumentParser( description="Create Live Install Media",
                                      fromfile_prefix_chars="@" )

    # These are mutually exclusive, one is required
    action = parser.add_mutually_exclusive_group( required=True )
    action.add_argument( "--make-iso", action="store_true",
                         help="Build a live iso" )
    action.add_argument( "--make-disk", action="store_true",
                         help="Build a partitioned disk image" )
    action.add_argument( "--make-fsimage", action="store_true",
                         help="Build a filesystem image" )
    action.add_argument( "--make-appliance", action="store_true",
                         help="Build an appliance image and XML description" )
    action.add_argument( "--make-ami", action="store_true",
                         help="Build an ami image" )
    action.add_argument( "--make-tar", action="store_true",
                         help="Build a tar of the root filesystem" )
    action.add_argument( "--make-pxe-live", action="store_true",
                         help="Build a live pxe boot squashfs image" )
    action.add_argument( "--make-ostree-live", action="store_true",
                         help="Build a live pxe boot squashfs image of Atomic Host" )


    parser.add_argument( "--iso", type=os.path.abspath,
                        help="Anaconda installation .iso path to use for virt-install" )
    parser.add_argument( "--fs-label", default="Anaconda",
                        help="Label to set on fsimage, default is 'Anaconda'")
    parser.add_argument("--compression", default="xz",
                        help="Compression binary for make-tar. xz, lzma, gzip, and bzip2 are supported. xz is the default.")
    parser.add_argument("--compress-arg", action="append", dest="compress_args", default=[],
                        help="Arguments to pass to compression. Pass once for each argument")

    parser.add_argument( "--ks", action="append", type=os.path.abspath,
                         help="Kickstart file defining the install." )
    parser.add_argument( "--image-only", action="store_true",
                         help="Exit after creating fs/disk image." )

    parser.add_argument( "--no-virt", action="store_true",
                         help="Use Anaconda's image install instead of virt-install" )
    parser.add_argument( "--proxy",
                         help="proxy URL to use for the install" )
    parser.add_argument( "--anaconda-arg", action="append", dest="anaconda_args",
                         help="Additional argument to pass to anaconda (no-virt "
                              "mode). Pass once for each argument" )
    parser.add_argument( "--armplatform",
                         help="the platform to use when creating images for ARM, "
                              "i.e., highbank, mvebu, omap, tegra, etc." )
    parser.add_argument( "--location", default=None, type=os.path.abspath,
                         help="location of iso directory tree with initrd.img "
                              "and vmlinuz. Used to run virt-install with a "
                              "newer initrd than the iso." )

    parser.add_argument( "--logfile", default="./livemedia.log",
                         type=os.path.abspath,
                         help="Path to logfile" )
    parser.add_argument( "--lorax-templates", default="/usr/share/lorax/",
                         type=os.path.abspath,
                         help="Path to mako templates for lorax" )
    parser.add_argument( "--tmp", default="/var/tmp", type=os.path.abspath,
                         help="Top level temporary directory" )
    parser.add_argument( "--resultdir", default=None, dest="result_dir",
                         type=os.path.abspath,
                         help="Directory to copy the resulting images and iso into. "
                              "Defaults to the temporary working directory")

    parser.add_argument( "--macboot", action="store_true",
                         dest="domacboot")
    parser.add_argument( "--nomacboot", action="store_false", default=False,
                         dest="domacboot")

    image_group = parser.add_argument_group("disk/fs image arguments")
    image_group.add_argument( "--disk-image", type=os.path.abspath,
                        help="Path to disk image to use for creating final image" )
    image_group.add_argument( "--keep-image", action="store_true",
                         help="Keep raw disk image after .iso creation" )
    image_group.add_argument( "--fs-image", type=os.path.abspath,
                        help="Path to filesystem image to use for creating final image" )
    image_group.add_argument( "--image-name", default=None,
                         help="Name of fs/disk image to create. Default is a random name." )
    image_group.add_argument("--qcow2", action="store_true",
                        help="create qcow2 image instead of raw sparse image")
    image_group.add_argument("--qcow2-arg", action="append", dest="qcow2_args", default=[],
                        help="Arguments to pass to qemu-img. Pass once for each argument")

    # Group of arguments for appliance creation
    app_group = parser.add_argument_group("appliance arguments")
    app_group.add_argument( "--app-name", default=None,
                            help="Name of appliance to pass to template")
    app_group.add_argument( "--app-template", default=None,
                         help="Path to template to use for appliance data.")
    app_group.add_argument( "--app-file", default="appliance.xml",
                         help="Appliance template results file.")

    # Group of arguments to pass to virt-install
    if not libvirt:
        virt_group = parser.add_argument_group("virt-install arguments (DISABLED -- no libvirt)")
    else:
        virt_group = parser.add_argument_group("virt-install arguments")
    virt_group.add_argument("--ram", metavar="MEMORY", type=int, default=2048,
                            help="Memory to allocate for installer in megabytes." )
    virt_group.add_argument("--vcpus", default=1,
                            help="Passed to --vcpus command" )
    virt_group.add_argument("--vnc",
                            help="Passed to --graphics command" )
    virt_group.add_argument("--arch", default=None,
                            help="Passed to --arch command")
    virt_group.add_argument("--kernel-args",
                            help="Additional argument to pass to the installation kernel")
    virt_group.add_argument("--ovmf-path", default="/usr/share/OVMF/",
                            help="Path to OVMF firmware. Requires OVMF_CODE.fd and OVMF_VARS.fd")
    virt_group.add_argument("--virt-uefi", action="store_true", default=False,
                            help="Use OVMF firmware to boot the VM in UEFI mode")

    # dracut arguments
    dracut_group = parser.add_argument_group( "dracut arguments" )
    dracut_group.add_argument( "--dracut-arg", action="append", dest="dracut_args",
                               help="Argument to pass to dracut when "
                                    "rebuilding the initramfs. Pass this "
                                    "once for each argument. NOTE: this "
                                    "overrides the default. (default: %s)" % (DRACUT_DEFAULT,) )

    # pxe to live arguments
    pxelive_group = parser.add_argument_group( "pxe to live arguments" )
    pxelive_group.add_argument( "--live-rootfs-size", type=int, default=0,
                                 help="Size of root filesystem of live image in GiB" )
    pxelive_group.add_argument( "--live-rootfs-keep-size", action="store_true",
                                help="Keep the original size of root filesystem in live image " )


    parser.add_argument( "--title", default="Linux Live Media",
                         help="Substituted for @TITLE@ in bootloader config files" )
    parser.add_argument( "--project", default="Red Hat Enterprise Linux",
                         help="substituted for @PROJECT@ in bootloader config files" )
    parser.add_argument( "--releasever", default="7",
                         help="substituted for @VERSION@ in bootloader config files" )
    parser.add_argument( "--volid", default=None, help="volume id")
    parser.add_argument( "--squashfs_args",
                         help="additional squashfs args" )

    # add the show version option
    parser.add_argument("-V", help="show program's version number and exit",
                        action="version", version=VERSION)

    return parser


def setup_logging(opts):
    # Setup logging to console and to logfile
    log.setLevel(logging.DEBUG)
    pylorax_log.setLevel(logging.DEBUG)

    sh = logging.StreamHandler()
    sh.setLevel(logging.INFO)
    fmt = logging.Formatter("%(asctime)s: %(message)s")
    sh.setFormatter(fmt)
    log.addHandler(sh)
    pylorax_log.addHandler(sh)

    fh = logging.FileHandler(filename=opts.logfile, mode="w")
    fh.setLevel(logging.DEBUG)
    fmt = logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s")
    fh.setFormatter(fmt)
    log.addHandler(fh)
    pylorax_log.addHandler(fh)

    # External program output log
    program_log.setLevel(logging.DEBUG)
    logfile = os.path.abspath(os.path.dirname(opts.logfile))+"/program.log"
    fh = logging.FileHandler(filename=logfile, mode="w")
    fh.setLevel(logging.DEBUG)
    program_log.addHandler(fh)


def default_image_name(compression, basename):
    """ Return a default image name with the correct suffix for the compression type.

    :param str compression: Compression type
    :param str basename: Base filename
    :returns: basename with compression suffix

    If the compression is unknown it defaults to xz
    """
    SUFFIXES = {"xz": ".xz", "gzip": ".gz", "bzip2": ".bz2", "lzma": ".lzma"}
    return basename + SUFFIXES.get(compression, ".xz")


def main():
    # parse the arguments
    parser = lorax_parser()
    opts = parser.parse_args()

    setup_logging(opts)

    log.info("livemedia-creator %s", vernum)
    log.debug( opts )

    if os.getuid() != 0:
        log.error("You need to run this as root")
        sys.exit( 1 )

    if opts.make_iso and not os.path.exists( opts.lorax_templates ):
        log.error( "The lorax templates directory (%s) doesn't"
                   " exist.", opts.lorax_templates)
        sys.exit( 1 )

    if opts.result_dir and os.path.exists(opts.result_dir):
        log.error( "The results_dir (%s) should not exist, please delete or "
                   "move its contents", opts.result_dir)
        sys.exit( 1 )

    # Default to putting results into tmp
    if not opts.result_dir:
        opts.result_dir = opts.tmp
    else:
        os.makedirs(opts.result_dir)

    if opts.iso and not os.path.exists( opts.iso ):
        log.error( "The iso %s is missing.", opts.iso)
        sys.exit( 1 )

    if opts.disk_image and not os.path.exists( opts.disk_image ):
        log.error( "The disk image %s is missing.", opts.disk_image)
        sys.exit( 1 )

    if opts.fs_image and not os.path.exists( opts.fs_image ):
        log.error( "The filesystem image %s is missing.", opts.fs_image)
        sys.exit( 1 )

    is_install = not (opts.disk_image or opts.fs_image)
    if is_install and not opts.no_virt and not opts.iso:
        log.error( "virt install needs an install iso." )
        sys.exit( 1 )

    if opts.volid and len(opts.volid) > 32:
        log.fatal("the volume id cannot be longer than 32 characters")
        sys.exit(1)

    if is_install and not opts.no_virt and not libvirt:
        log.error("virt install requires libvirt-python to be installed.")
        sys.exit(1)

    if is_install and not opts.no_virt \
       and not os.path.exists("/usr/bin/virt-install"):
        log.error("virt-install needs to be installed.")
        sys.exit(1)

    if is_install and opts.no_virt \
       and not os.path.exists("/usr/sbin/anaconda"):
        log.error("no-virt requires anaconda to be installed.")
        sys.exit(1)

    if opts.make_appliance and not opts.app_template:
        opts.app_template = joinpaths(opts.lorax_templates,
                                            "appliance/libvirt.tmpl")

    if opts.make_appliance and not os.path.exists(opts.app_template):
        log.error("The appliance template (%s) doesn't "
                  "exist", opts.app_template)
        sys.exit(1)

    if opts.image_name and os.path.exists(joinpaths(opts.result_dir, opts.image_name)):
        log.error("The disk image to be created should not exist.")
        sys.exit(1)

    if opts.qcow2 and not os.path.exists("/usr/bin/qemu-img"):
        log.error("qcow2 requires the qemu-img utility to be installed.")
        sys.exit(1)

    if opts.qcow2 and opts.make_iso:
        log.error("qcow2 cannot be used to make a bootable iso.")
        sys.exit(1)

    if opts.virt_uefi:
        if not os.path.isdir(opts.ovmf_path):
            log.error("The OVMF firmware is missing from %s", opts.ovmf_path)
            sys.exit(1)

        for f in ["OVMF_CODE.fd", "OVMF_VARS.fd"]:
            if not os.path.exists(joinpaths(opts.ovmf_path, f)):
                log.error("OVMF firmware file %s is missing from %s", f, opts.ovmf_path)
                sys.exit(1)

    # AMI image is just a fsimage with an AMI label
    if opts.make_ami:
        opts.make_fsimage = True
        if not opts.image_name:
            opts.image_name = "ami-root.img"
        if opts.fs_label == "Anaconda":
            opts.fs_label = "AMI"
    elif opts.make_tar:
        if not opts.image_name:
            opts.image_name = default_image_name(opts.compression, "root.tar")
        if opts.compression == "xz" and not opts.compress_args:
            opts.compress_args = ["-9"]

    if opts.app_file:
        opts.app_file = joinpaths(opts.result_dir, opts.app_file)

    if opts.make_ostree_live:
        opts.make_pxe_live = True
        opts.ostree = True
    else:
        opts.ostree = False

    tempfile.tempdir = opts.tmp
    disk_img = None

    # Parse the kickstart
    if opts.ks:
        ks_version = makeVersion(RHEL7)
        ks = KickstartParser( ks_version, errorsAreFatal=False, missingIncludeIsFatal=False )
        ks.readKickstart( opts.ks[0] )

    # Make the disk or filesystem image
    if not opts.disk_image and not opts.fs_image:
        if not opts.ks:
            log.error("Image creation requires a kickstart file")
            sys.exit(1)

        errors = []
        if ks.handler.method.method != "url" and opts.no_virt:
            errors.append("Only url install method is currently supported. Please "
                          "fix your kickstart file." )

        if ks.handler.displaymode.displayMode is not None:
            errors.append("The kickstart must not set a display mode (text, cmdline, "
                          "graphical), this will interfere with livemedia-creator.")

        if opts.make_fsimage:
            # Make sure the kickstart isn't using autopart and only has a / mountpoint
            part_ok = not any(p for p in ks.handler.partition.partitions
                                 if p.mountpoint not in ["/", "swap"])
            if not part_ok or ks.handler.autopart.seen:
                errors.append("Filesystem images must use a single / part, not autopart or "
                              "multiple partitions. swap is allowed but not used.")

        if errors:
            for _e in errors:
                log.error(_e)
            sys.exit(1)

        # Make the image. Output of this is either a partitioned disk image or a fsimage
        try:
            disk_img = make_image(opts, ks)
        except InstallError as e:
            log.error("ERROR: Image creation failed: %s", e)
            sys.exit(1)

    if not opts.image_only:
        result_dir = None
        if opts.make_iso:
            work_dir = tempfile.mkdtemp()
            log.info("working dir is %s", work_dir)

            if (opts.fs_image or opts.no_virt) and not opts.disk_image:
                # Create iso from a filesystem image
                disk_img = opts.fs_image or disk_img

                make_squashfs(disk_img, work_dir)
                with Mount(disk_img, opts="loop") as mount_dir:
                    result_dir = make_livecd(opts, mount_dir, work_dir)
            else:
                # Create iso from a partitioned disk image
                disk_img = opts.disk_image or disk_img
                with PartitionMount(disk_img) as img_mount:
                    if img_mount and img_mount.mount_dir:
                        make_runtime(opts, img_mount.mount_dir, work_dir)
                        result_dir = make_livecd(opts, img_mount.mount_dir, work_dir)

            # cleanup the mess
            # cleanup work_dir?
            if disk_img and not (opts.keep_image or opts.disk_image or opts.fs_image):
                os.unlink(disk_img)
                log.info("Disk image erased")
                disk_img = None
        elif opts.make_appliance:
            if not opts.ks:
                networks = []
            else:
                networks = ks.handler.network.network
            make_appliance(opts.disk_image or disk_img, opts.app_name,
                           opts.app_template, opts.app_file, networks, opts.ram,
                           opts.vcpus, opts.arch, opts.title, opts.project, opts.releasever)
        elif opts.make_pxe_live:
            work_dir = tempfile.mkdtemp()
            log.info("working dir is %s", work_dir)

            if (opts.fs_image or opts.no_virt) and not opts.disk_image:
                # Create pxe live images from a filesystem image
                disk_img = opts.fs_image or disk_img
                with Mount(disk_img, opts="loop") as mnt_dir:
                    result_dir = make_live_images(opts, work_dir, mnt_dir, rootfs_image=disk_img)
            else:
                # Create pxe live images from a partitioned disk image
                disk_img = opts.disk_image or disk_img
                is_root_part = None
                if opts.ostree:
                    is_root_part = lambda dir: os.path.exists(dir+"/ostree/deploy")
                with PartitionMount(disk_img, mount_ok=is_root_part) as img_mount:
                    if img_mount and img_mount.mount_dir:
                        try:
                            mounted_sysroot_boot_dir = None
                            if opts.ostree:
                                mounted_sysroot_boot_dir = mount_boot_part_over_root(img_mount)
                            if opts.live_rootfs_keep_size:
                                size = img_mount.mount_size / 1024**3
                            else:
                                size = opts.live_rootfs_size or None
                            result_dir = make_live_images(opts, work_dir, img_mount.mount_dir, size=size)
                        finally:
                            if mounted_sysroot_boot_dir:
                                umount(mounted_sysroot_boot_dir)

        if opts.result_dir != opts.tmp and result_dir:
            copytree(result_dir, opts.result_dir, preserve=False)
            shutil.rmtree( result_dir )

    log.info("SUMMARY")
    log.info("-------")
    log.info("Logs are in %s", os.path.abspath(os.path.dirname(opts.logfile)))
    if disk_img:
        log.info("Disk image is at %s", disk_img)
    if opts.make_appliance:
        log.info("Appliance description is in %s", opts.app_file)
    log.info("Results are in %s", opts.result_dir or result_dir)

    sys.exit( 0 )

if __name__ == '__main__':
    main()