#!/usr/bin/python3
#
# kstest-runner - LMC-like program that only does what's necessary for running
# kickstart-based tests in anaconda
#
# Copyright (C) 2011-2015  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>
#            Chris Lumens <clumens@redhat.com>
#
import logging
log = logging.getLogger("livemedia-creator")

import os
import sys
import uuid
import tempfile
import subprocess
from time import sleep
import argparse

# Use the Lorax treebuilder branch for iso creation
from pylorax.treebuilder import udev_escape
from pylorax.executils import execWithRedirect

from pylorax import setup_logging
from pylorax.monitor import LogMonitor
from pylorax.mount import IsoMountpoint

import libvirt

class InstallError(Exception):
    pass

class VirtualInstall(object):
    """
    Run virt-install using an iso and a kickstart
    """
    def __init__(self, iso, ks_paths, disk_paths,
                 kernel_args=None, memory=1024, vnc=None,
                 log_check=None, virtio_host="127.0.0.1", virtio_port=6080):
        """
        Start the installation

        :param iso: Information about the iso to use for the installation
        :type iso: IsoMountpoint
        :param list ks_paths: Paths to kickstart files. All are injected, the
           first one is the one executed.
        :param list disk_paths: Paths to pre-existing disk images.
        :param str kernel_args: Extra kernel arguments to pass on the kernel cmdline
        :param int memory: Amount of RAM to assign to the virt, in MiB
        :param str vnc: Arguments to pass to virt-install --graphics
        :param log_check: Method that returns True if the installation fails
        :type log_check: method
        :param str virtio_host: Hostname to connect virtio log to
        :param int virtio_port: Port to connect virtio log to
        """
        self.virt_name = "LiveOS-"+str(uuid.uuid4())
        # add --graphics none later
        # add whatever serial cmds are needed later
        args = ["-n", self.virt_name,
                 "-r", str(memory),
                 "--noreboot",
                 "--noautoconsole"]

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

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

        for disk in disk_paths:
            args.append("--disk")
            args.append("path={0}".format(disk))

        if iso.stage2:
            disk_opts = "path={0},device=cdrom,readonly=on,shareable=on".format(iso.iso_path)
            args.append("--disk")
            args.append(disk_opts)

        extra_args = "ks=file:/{0}".format(os.path.basename(ks_paths[0]))
        if not vnc:
            extra_args += " inst.cmdline"
        if kernel_args:
            extra_args += " " + kernel_args
        if iso.stage2:
            extra_args += " stage2=hd:CDLABEL={0}".format(udev_escape(iso.label))
        args.append("--extra-args")
        args.append(extra_args)

        args.append("--location")
        args.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)
        args.append("--channel")
        args.append(channel_args)

        log.info("Running virt-install.")
        try:
            execWithRedirect("virt-install", args, raise_err=True)
        except subprocess.CalledProcessError as e:
            raise InstallError("Problem starting virtual install: %s" % e)

        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, poolName):
        """
        Make sure the virt has been shut down and destroyed

        Could use libvirt for this instead.
        """
        log.info("Shutting down %s", self.virt_name)
        subprocess.call(["virsh", "destroy", self.virt_name])
        subprocess.call(["virsh", "undefine", self.virt_name])
        subprocess.call(["virsh", "pool-destroy", poolName])
        subprocess.call(["virsh", "pool-undefine", poolName])

def virt_install(opts, install_log):
    """
    Use virt-install to install to a disk image

    :param opts: options passed to livemedia-creator
    :type opts: argparse options
    :param str install_log: The path to write the log from virt-install

    This uses virt-install with a boot.iso and a kickstart to create a disk
    image.
    """
    iso_mount = IsoMountpoint(opts.iso, opts.location)
    log_monitor = LogMonitor(install_log, timeout=opts.timeout)

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

    try:
        virt = VirtualInstall(iso_mount, opts.ks, opts.disk_paths,
                              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(os.path.basename(opts.tmp))
        log_monitor.shutdown()
    except InstallError as e:
        log.error("VirtualInstall failed: %s", e)
        raise
    finally:
        log.info("unmounting the iso")
        iso_mount.umount()

    if log_monitor.server.log_check():
        if not log_monitor.server.error_line and opts.timeout:
            msg = "virt_install failed due to timeout"
        else:
            msg = "virt_install failed on line: %s" % log_monitor.server.error_line
        raise InstallError(msg)


def make_image(opts):
    """
    Install to a disk image

    :param opts: options passed to livemedia-creator
    :type opts: argparse options
    :param str ks: Path to the kickstart to use for the installation

    Use virt-install or anaconda to install to a disk image.
    """
    try:
        install_log = os.path.abspath(os.path.dirname(opts.logfile))+"/virt-install.log"
        log.info("install_log = %s", install_log)

        virt_install(opts, install_log)
    except InstallError as e:
        log.error("Install failed: %s", e)

        if not opts.keep_image:
            log.info("Removing bad disk image")

            for image in opts.disk_paths:
                if os.path.exists(image):
                    os.unlink(image)

        raise

    log.info("Disk Image install successful")

def main():
    parser = argparse.ArgumentParser(description="Run anaconda's kickstart-based tests",
                                     fromfile_prefix_chars="@")

    parser.add_argument("--iso", type=os.path.abspath,
                        help="Anaconda installation .iso path to use for virt-install")
    parser.add_argument("--ks", action="append", type=os.path.abspath,
                        help="Kickstart file defining the install.")
    parser.add_argument("--disk", action="append", type=os.path.abspath, dest="disk_paths",
                        help="Pre-existing disk image to use for destination.")

    parser.add_argument("--proxy",
                        help="proxy URL to use for the install")
    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("--tmp", default="/var/tmp", type=os.path.abspath,
                        help="Top level temporary directory")

    image_group = parser.add_argument_group("disk/fs image arguments")
    image_group.add_argument("--keep-image", action="store_true",
                             help="Keep raw disk image after .iso creation")

    # Group of arguments to pass to virt-install
    virt_group = parser.add_argument_group("virt-install arguments")
    virt_group.add_argument("--ram", metavar="MEMORY", type=int, default=1024,
                            help="Memory to allocate for installer in megabytes.")
    virt_group.add_argument("--vnc",
                            help="Passed to --graphics command")
    virt_group.add_argument("--kernel-args",
                            help="Additional argument to pass to the installation kernel")

    parser.add_argument("--timeout", default=None, type=int,
                        help="Cancel installer after X minutes")

    opts = parser.parse_args()

    setup_logging(opts.logfile, log)

    log.debug( opts )

    # Check for invalid combinations of options, print all the errors and exit.
    errors = []
    if opts.ks and not os.path.exists(opts.ks[0]):
        errors.append("kickstart file (%s) is missing." % opts.ks[0])

    if opts.iso and not os.path.exists(opts.iso):
        errors.append("The iso %s is missing." % opts.iso)

    if not opts.iso:
        errors.append("virt install needs an install iso.")

    if not os.path.exists("/usr/bin/virt-install"):
        errors.append("virt-install needs to be installed.")

    if os.getuid() != 0:
        errors.append("You need to run this as root")

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

    tempfile.tempdir = opts.tmp

    # Make the image.
    try:
        make_image(opts)
    except InstallError as e:
        log.error("ERROR: Image creation failed: %s", e)
        sys.exit(1)

    log.info("SUMMARY")
    log.info("-------")
    log.info("Logs are in %s", os.path.abspath(os.path.dirname(opts.logfile)))
    log.info("Disk image(s) at %s", ",".join(opts.disk_paths))
    log.info("Results are in %s", opts.tmp)

    sys.exit( 0 )

if __name__ == '__main__':
    main()