lorax/src/sbin/livemedia-creator

#!/usr/bin/env python
#
# Live Media Creator
#
# Copyright (C) 2011-2012  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/>.
#
# Author(s): Brian C. Lane <bcl@redhat.com>
#
import logging
log = logging.getLogger("livemedia-creator")
program_log = logging.getLogger("program")
pylorax_log = logging.getLogger("pylorax")

import os
import sys
import uuid
import tempfile
import subprocess
import socket
import threading
import SocketServer
import libvirt
from time import sleep
import shutil
import traceback
import argparse

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

# Use the Lorax treebuilder branch for iso creation
from pylorax.base import DataHolder
from pylorax.treebuilder import TreeBuilder, RuntimeBuilder, udev_escape
from pylorax.sysutils import joinpaths, remove, linktree
from pylorax.imgutils import PartitionMount, mksparse
from pylorax.executils import execWithRedirect, execWithCapture


# Default parameters for rebuilding initramfs, overrice with --dracut-args
DRACUT_DEFAULT = ["--xz", "--add", "livenet", "--add", "dmsquash-live",
                  "--omit", "plymouth"]


class LogRequestHandler(SocketServer.BaseRequestHandler):
    """
    Handle monitoring and saving the logfiles from the virtual install
    """
    def setup(self):
        if self.server.log_path:
            self.fp = open(self.server.log_path, "w")
        else:
            print "no log_path specified"
        self.request.settimeout(10)

    def handle(self):
        """
        Handle writing incoming data to a logfile and
        checking the logs for any Tracebacks or other errors that indicate
        that the install failed.
        """
        line = ""
        while True:
            if self.server.kill:
                break

            try:
                data = self.request.recv(4096)
                self.fp.write(data)
                self.fp.flush()

                # check the data for errors and set error flag
                # need to assemble it into lines so we can test for the error
                # string.
                while data:
                    more = data.split("\n", 1)
                    line += more[0]
                    if len(more) > 1:
                        self.iserror(line)
                        line = ""
                        data = more[1]
                    else:
                        data = None

            except socket.timeout:
                pass
            except:
                break

    def finish(self):
        self.fp.close()

    def iserror(self, line):
        """
        Check a line to see if it contains an error indicating install failure
        """
        simple_tests = [ "Traceback (",
                         "Out of memory:",
                         "Call Trace:",
                         "insufficient disk space:" ]
        for t in simple_tests:
            if line.find( t ) > -1:
                self.server.log_error = True
                return


class LogServer(SocketServer.TCPServer):
    """
    Add path to logfile
    Add log error flag
    Add a kill switch
    """
    def __init__(self, log_path, *args, **kwargs):
        self.kill = False
        self.log_error = False
        self.log_path = log_path
        SocketServer.TCPServer.__init__(self, *args, **kwargs)

    def log_check(self):
        return self.log_error


class LogMonitor(object):
    """
    Contains all the stuff needed to setup a thread to listen to the logs
    from the virtual install
    """
    def __init__(self, log_path, host="localhost", port=0):
        """
        Fire up the thread listening for logs
        """
        self.server = LogServer(log_path, (host, port), LogRequestHandler)
        self.host, self.port = self.server.server_address
        self.log_path = log_path
        self.server_thread = threading.Thread(target=self.server.handle_request)
        self.server_thread.daemon = True
        self.server_thread.start()

    def shutdown(self):
        self.server.kill = True
        self.server_thread.join()


class IsoMountpoint(object):
    """
    Mount the iso on a temporary directory and check to make sure the
    vmlinuz and initrd.img files exist
    Check the iso for a LiveOS directory and set a flag.
    Extract the iso's label.
    """
    def __init__( self, iso_path ):
        self.label = None
        self.iso_path = iso_path
        self.mount_dir = tempfile.mkdtemp()
        if not self.mount_dir:
            raise Exception("Error creating temporary directory")

        cmd = ["mount","-o","loop",self.iso_path,self.mount_dir]
        log.debug( cmd )
        execWithRedirect( cmd[0], cmd[1:] )

        self.kernel = self.mount_dir+"/isolinux/vmlinuz"
        self.initrd = self.mount_dir+"/isolinux/initrd.img"

        if os.path.isdir( self.mount_dir+"/repodata" ):
            self.repo = self.mount_dir
        else:
            self.repo = None
        self.liveos = os.path.isdir( self.mount_dir+"/LiveOS" )

        try:
            for f in [self.kernel, self.initrd]:
                if not os.path.isfile(f):
                    raise Exception("Missing file on iso: {0}".format(f))
        except:
            self.umount()
            raise

        self.get_iso_label()

    def umount( self ):
        cmd = ["umount", self.mount_dir]
        log.debug( cmd )
        execWithRedirect( cmd[0], cmd[1:] )
        os.rmdir( self.mount_dir )

    def get_iso_label( self ):
        """
        Get the iso's label using isoinfo
        """
        cmd = [ "isoinfo", "-d", "-i", self.iso_path ]
        log.debug( cmd )
        isoinfo_output = execWithCapture( cmd[0], cmd[1:] )
        log.debug( isoinfo_output )
        for line in isoinfo_output.splitlines():
            if line.startswith("Volume id: "):
                self.label = line[11:]
                return


class ImageMount(object):
    """
    # kpartx -p p -v -a /tmp/diskV2DiCW.im
    # add map loop2p1 (253:2): 0 3481600 linear /dev/loop2 2048
    # add map loop2p2 (253:3): 0 614400 linear /dev/loop2 3483648

    Find the / partition and only mount that
    """
    def __init__(self, disk_img):
        """
        Use kpartx to mount an image onto temporary directories
        return a list of the dirs
        """
        self.disk_img = disk_img

        # call kpartx and parse the output
        cmd = [ "kpartx", "-v", "-p", "p", "-a", disk_img ]
        log.debug( cmd )
        kpartx_output = execWithCapture( cmd[0], cmd[1:] )
        log.debug( kpartx_output )
        self.loop_devices = []
        for line in kpartx_output.splitlines():
            # add map loop2p3 (253:4): 0 7139328 linear /dev/loop2 528384
            # 3rd element is size in 512 byte blocks
            if line.startswith("add map "):
                fields = line[8:].split()
                self.loop_devices.append((fields[0], int(fields[3])*512))

        log.debug( self.loop_devices )

        # Mount the devices, if possible
        self.mount_dir = None
        mount_dir = tempfile.mkdtemp()
        for dev, size in self.loop_devices:
            cmd = ["mount", "/dev/mapper/"+dev, mount_dir]
            log.debug( cmd )
            try:
                execWithRedirect( cmd[0], cmd[1:] )
                if os.path.isfile(mount_dir+"/etc/fstab"):
                    self.mount_dir = mount_dir
                    self.mount_dev = dev
                    self.mount_size = size
                    break
                cmd = ["umount", mount_dir]
                execWithRedirect( cmd[0], cmd[1:] )
            except subprocess.CalledProcessError:
                log.debug( traceback.format_exc() )

        if self.mount_dir:
            log.info( "Found root partition, mounted on {0}".format(self.mount_dir) )
            log.info( "size = {0}".format(self.mount_size) )

    def umount(self):
        """
        unmount the disk image
        """
        if self.mount_dir:
            cmd = ["umount", self.mount_dir]
            log.debug( cmd )
            execWithRedirect( cmd[0], cmd[1:] )
            os.rmdir(self.mount_dir)
            self.mount_dir = None

        cmd = ["kpartx", "-d", self.disk_img]
        log.debug( cmd )
        execWithRedirect( cmd[0], cmd[1:] )


class VirtualInstall( object ):
    """
    Run virt-install using an iso and kickstart(s)
    """
    def __init__( self, iso, ks_paths, disk_img, img_size=2, 
                  kernel_args=None, memory=1024, vnc=None,
                  log_check=None, virtio_host="127.0.0.1", virtio_port=6080 ):
        """

        iso is an instance of IsoMountpoint
        ks_paths is a list of paths to a kickstart files. All are injected, the
                 first one is the one executed.
        disk_img is the path to a disk image (doesn't need to exist)
        img_size is the size, in GiB, of the image if it doesn't exist
        kernel_args are extra arguments to pass on the kernel cmdline
        memory is the amount of ram to assign to the virt
        vnc is passed to the --graphics command verbatim
        log_check is a method that returns True of the log indicates an error
        virtio_host and virtio_port are used to communicate with the log monitor
        """
        self.virt_name = "LiveOS-"+str(uuid.uuid4())
        # add --graphics none later
        # add whatever serial cmds are needed later
        cmd = [ "virt-install", "-n", self.virt_name,
                                "-r", str(memory),
                                "--noreboot",
                                "--noautoconsole" ]

        cmd.append("--graphics")
        if vnc:
            cmd.append(vnc)
        else:
            cmd.append("none")

        for ks in ks_paths:
            cmd.append("--initrd-inject")
            cmd.append(ks)

        disk_opts = "path={0}".format(disk_img)
        disk_opts += ",format=raw"
        if not os.path.isfile(disk_img):
            disk_opts += ",size={0}".format(img_size)
        cmd.append("--disk")
        cmd.append(disk_opts)

        if iso.liveos:
            disk_opts = "path={0},device=cdrom".format(iso.iso_path)
            cmd.append("--disk")
            cmd.append(disk_opts)

        extra_args = "ks=file:/{0}".format(os.path.basename(ks_paths[0]))
        if kernel_args:
            extra_args += " "+kernel_args
        if not vnc:
            extra_args += " console=/dev/ttyS0"
        if iso.liveos:
            extra_args += " root=live:CDLABEL={0}".format(udev_escape(iso.label))
        cmd.append("--extra-args")
        cmd.append(extra_args)

        cmd.append("--location")
        cmd.append(iso.mount_dir)

        channel_args = "tcp,host={0}:{1},mode=connect,target_type=virtio" \
                       ",name=org.fedoraproject.anaconda.log.0".format(
                       virtio_host, virtio_port)
        cmd.append("--channel")
        cmd.append(channel_args)

        log.debug( cmd )

        rc = execWithRedirect( cmd[0], cmd[1:] )
        if rc:
            raise Exception("Problem starting virtual install")

        conn = libvirt.openReadOnly(None)
        dom = conn.lookupByName(self.virt_name)

        # TODO: If vnc has been passed, we should look up the port and print that
        # for the user at this point

        while dom.isActive() and not log_check():
            sys.stdout.write(".")
            sys.stdout.flush()
            sleep(10)
        print

        if log_check():
            log.info( "Installation error detected. See logfile." )
        else:
            log.info( "Install finished. Or at least virt shut down." )

    def destroy( self ):
        """
        Make sure the virt has been shut down and destroyed

        Could use libvirt for this instead.
        """
        log.info( "Shutting down {0}".format(self.virt_name) )
        subprocess.call(["virsh","destroy",self.virt_name])
        subprocess.call(["virsh","undefine",self.virt_name])


def anaconda_install( disk_img, disk_size, kickstart, repo, args ):
    """

    """
    # Create the sparse image
    mksparse( disk_img, disk_size * 1024**3 )

    cmd = [ "anaconda", "--image", disk_img, "--kickstart", kickstart,
            "--script", "--repo", repo_url ]
    cmd += args

    log.debug( cmd )

    return execWithRedirect( cmd[0], cmd[1:] )


def get_kernels( boot_dir ):
    """
    Examine the vmlinuz-* versions and return a list of them
    """
    files = os.listdir(boot_dir)
    return [f[8:] for f in files if f.startswith("vmlinuz-")]


def make_livecd( disk_img, squashfs_args="", templatedir=None,
                 title="Linux", project="Linux", releasever=16 ):
    """
    Take the content from the disk image and make a livecd out of it

    This uses wwood's squashfs live initramfs method:
     * put the real / into LiveOS/rootfs.img
     * make a squashfs of the LiveOS/rootfs.img tree
     * make a simple initramfs with the squashfs.img and /etc/cmdline in it
     * make a cpio of that tree
     * append the squashfs.cpio to a dracut initramfs for each kernel installed

    Then on boot dracut reads /etc/cmdline which points to the squashfs.img
    mounts that and then mounts LiveOS/rootfs.img as /

    """
    with PartitionMount( disk_img ) as img_mount:
        kernel_list = get_kernels( joinpaths( img_mount.mount_dir, "boot" ) )
        log.debug( "kernel_list = {0}".format(kernel_list) )
        if kernel_list:
            kernel_arch = kernel_list[0].split(".")[-1]
        else:
            kernel_arch = "i386"
        log.debug( "kernel_arch = {0}".format(kernel_arch) )

        work_dir = tempfile.mkdtemp()
        log.info( "working dir is {0}".format(work_dir) )

        runtime = "images/install.img"
        # This is a mounted image partition, cannot hardlink to it, so just use it
        installroot = img_mount.mount_dir
        # symlink installroot/images to work_dir/images so we don't run out of space
        os.makedirs( joinpaths( work_dir, "images" ) )

        # TreeBuilder expects the config files to be in /tmp/config_files
        # I think these should be release specific, not from lorax, but for now
        configdir = joinpaths(templatedir,"config_files/live/")
        configdir_path = "tmp/config_files"
        fullpath = joinpaths(installroot, configdir_path)
        if os.path.exists(fullpath):
            remove(fullpath)
        shutil.copytree(configdir, fullpath)

        # Fake yum object
        fake_yum = DataHolder( conf=DataHolder( installroot=installroot ) )
        # Fake arch with only basearch set
        arch = DataHolder( basearch=kernel_arch, libdir=None, buildarch=None )
        # TODO: Need to get release info from someplace...
        product = DataHolder( name=project, version=releasever, release="",
                                variant="", bugurl="", isfinal=False )

        rb = RuntimeBuilder( product, arch, fake_yum )
        log.info( "Creating runtime" )
        rb.create_runtime( joinpaths( work_dir, runtime ), size=None )

        # Link /images to work_dir/images to make the templates happy
        if os.path.islink( joinpaths( installroot, "images" ) ):
            os.unlink( joinpaths( installroot, "images" ) )
        subprocess.check_call(["/bin/ln", "-s", joinpaths( work_dir, "images" ),
                                                joinpaths( installroot, "images" )])

        tb = TreeBuilder( product=product, arch=arch,
                          inroot=installroot, outroot=work_dir,
                          runtime=runtime, templatedir=templatedir)
        log.info( "Rebuilding initrds" )
        if not opts.dracut_args:
            dracut_args = DRACUT_DEFAULT
        else:
            dracut_args = []
            for arg in opts.dracut_args:
                dracut_args += arg.split(" ", 1)
        log.info( "dracut args = {0}".format( dracut_args ) )
        tb.rebuild_initrds( add_args=dracut_args )
        log.info( "Building boot.iso" )
        tb.build()

        return work_dir


if __name__ == '__main__':
    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 disk 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" )

    parser.add_argument( "--iso", type=os.path.abspath,
                        help="Anaconda installation .iso path to use for virt-install" )
    parser.add_argument( "--disk-image", type=os.path.abspath,
                        help="Path to disk image to use for creating final image" )

    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 disk image." )
    parser.add_argument( "--keep-image", action="store_true",
                         help="Keep raw disk image after .iso creation" )
    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( "--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="/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")

    # Group of arguments to pass to virt-install
    virt_group = parser.add_argument_group("virt-install arguments")
    virt_group.add_argument("--ram", metavar="MEMORY", default=1024,
                            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",
                            help="Passed to --arch command" )
    virt_group.add_argument( "--kernel-args",
                             help="Additional argument to pass to the installation kernel" )

    # 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,) )

    parser.add_argument( "--title", default="Linux Live Media",
                         help="Substituted for @TITLE@ in bootloader config files" )
    parser.add_argument( "--project", default="Linux",
                         help="substituted for @PROJECT@ in bootloader config files" )
    parser.add_argument( "--releasever", type=int, default=16,
                         help="substituted for @VERSION@ in bootloader config files" )
    parser.add_argument( "--squashfs_args",
                         help="additional squashfs args" )

    opts = parser.parse_args()

    # 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 )

    log.debug( opts )

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

    if not os.path.exists( opts.lorax_templates ):
        log.error( "The lorax templates directory ({0}) doesn't"
                   " exist.".format( opts.lorax_templates ) )
        sys.exit( 1 )

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

    if opts.iso and not os.path.exists( opts.iso ):
        log.error( "The iso {0} is missing.".format( opts.iso ) )
        sys.exit( 1 )

    if opts.disk_image and not os.path.exists( opts.disk_image ):
        log.error( "The disk image {0} is missing.".format( opts.disk_image ) )
        sys.exit( 1 )

    if opts.make_appliance:
        log.error( "--make-appliance is not yet implemented." )
        sys.exit( 1 )

    if opts.make_ami:
        log.error( "--make-ami is not yet implemented." )
        sys.exit( 1 )

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

    # Make the disk image
    if not opts.disk_image:
        # Parse the kickstart to get the partition sizes
        ks_version = makeVersion()
        ks = KickstartParser( ks_version, errorsAreFatal=False, missingIncludeIsFatal=False )
        ks.readKickstart( opts.ks[0] )
        disk_size = 1 + (sum( [p.size for p in ks.handler.partition.partitions] ) / 1024)
        log.info( "disk_size = {0}GB".format(disk_size) )

        if ks.handler.method.method != "url":
            log.error( "Only url install method is currently supported. Please "
                       "fix your kickstart file." )
            sys.exit( 1 )
        repo_url = ks.handler.method.url

        disk_img = tempfile.mktemp( prefix="disk", suffix=".img", dir=opts.tmp )
        install_log = os.path.abspath(os.path.dirname(opts.logfile))+"/virt-install.log"

        log.info( "disk_img = {0}".format(disk_img) )
        log.info( "install_log = {0}".format(install_log) )

        if opts.no_virt:
            anaconda_args = []
            if opts.anaconda_args:
                for arg in opts.anaconda_args:
                    anaconda_args += arg.split(" ", 1)
            if opts.proxy:
                anaconda_args += [ "--proxy", opts.proxy ]

            # Use anaconda's image install
            install_error = anaconda_install( disk_img, disk_size, opts.ks[0],
                                              repo_url, anaconda_args )

            # Move the anaconda logs over to a log directory
            log_dir = os.path.abspath(os.path.dirname(opts.logfile))
            log_anaconda = joinpaths( log_dir, "anaconda" )
            if not os.path.isdir( log_anaconda ):
                os.mkdir( log_anaconda )
            for l in ["anaconda.log", "ifcfg.log", "program.log", "storage.log",
                      "yum.log"]:
                if os.path.exists( "/tmp/"+l ):
                    shutil.copy2( "/tmp/"+l, log_anaconda )
                    os.unlink( "/tmp/"+l )
        else:
            iso_mount = IsoMountpoint( opts.iso )
            log_monitor = LogMonitor( install_log )

            kernel_args = ""
            if opts.kernel_args:
                kernel_args += opts.kernel_args
            if opts.proxy:
                kernel_args += " proxy="+opts.proxy

            virt = VirtualInstall( iso_mount, opts.ks, disk_img, disk_size,
                                   kernel_args, opts.ram, opts.vnc,
                                   log_check = log_monitor.server.log_check,
                                   virtio_host = log_monitor.host,
                                   virtio_port = log_monitor.port )
            virt.destroy()
            log_monitor.shutdown()
            iso_mount.umount()

            install_error = log_monitor.server.log_check()

        if install_error:
            log.error( "Install failed" )
            if not opts.keep_image:
                log.info( "Removing bad disk image" )
                os.unlink( disk_img )
            sys.exit( 1 )
        else:
            log.info( "Disk Image install successful" )


    result_dir = None
    if opts.make_iso and not opts.image_only:
        result_dir = make_livecd( opts.disk_image or disk_img, opts.squashfs_args,
                                  opts.lorax_templates,
                                  opts.title, opts.project, opts.releasever )

        if not opts.keep_image and not opts.disk_image:
            os.unlink( disk_img )

        if opts.result_dir:
            shutil.copytree( result_dir, opts.result_dir )
            shutil.rmtree( result_dir )

    log.info("SUMMARY")
    log.info("-------")
    log.info("Logs are in {0}".format(os.path.abspath(os.path.dirname(opts.logfile))))
    if opts.keep_image or opts.make_disk:
        log.info("Disk image is at {0}".format(disk_img))
    else:
        log.info("Disk image erased")
    if result_dir:
        log.info("Results are in {0}".format(opts.result_dir or result_dir))

    sys.exit( 0 )