#!/usr/bin/python -tt
#
# Script to set up a Xen guest and kick off an install
#
# Copyright 2005-2006  Red Hat, Inc.
# Jeremy Katz <katzj@redhat.com>
# Option handling added by Andrew Puch <apuch@redhat.com>
#
# This software may be freely redistributed under the terms of the GNU
# general public license.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.


import os, sys, string
from optparse import OptionParser, OptionValueError
import subprocess
import struct
import logging
import libxml2
import urlgrabber.progress as progress

import libvirt
import virtinst

MIN_RAM = 256

### Utility functions
def yes_or_no(s):
    s = s.lower()
    if s in ("y", "yes", "1", "true", "t"):
        return True
    elif s in ("n", "no", "0", "false", "f"):
        return False
    raise ValueError, "A yes or no response is required" 

def prompt_for_input(prompt = "", val = None):
    if val is not None:
        return val
    print prompt + " ", 
    return sys.stdin.readline().strip()


### General input gathering functions
def get_full_virt():
    while 1:
        res = prompt_for_input("Would you like a fully virtualized guest (yes or no)?  This will allow you to run unmodified operating systems.")
        try:
            return yes_or_no(res)
        except ValueError, e:
            print "ERROR: ", e

def get_name(name, guest):
    while 1:
        name = prompt_for_input("What is the name of your virtual machine?", name)
        try:
            guest.name = name
            break
        except ValueError, e:
            print "ERROR: ", e            
            name = None

def get_memory(memory, guest):
    while 1:
        memory = prompt_for_input("How much RAM should be allocated (in megabytes)?", memory)
        if memory < MIN_RAM:
            print "ERROR: Installs currently require %d megs of RAM." %(MIN_RAM,)
            print ""
            memory = None
            continue
        try:
            guest.memory = int(memory)
            break
        except ValueError, e:
            print "ERROR: ", e
            memory = None

def get_uuid(uuid, guest):
    if uuid: 
        try:
            guest.uuid = uuid
        except ValueError, e:
            print "ERROR: ", e
            sys.exit(1)

def get_vcpus(vcpus, check_cpu, guest, conn):
    while 1:
        if check_cpu is None:
            break
        hostinfo = conn.getInfo()
        cpu_num = hostinfo[4] * hostinfo[5] * hostinfo[6] * hostinfo[7]
        if vcpus <= cpu_num:
            break
        res = prompt_for_input("You have asked for more virtual CPUs (%d) than there are physical CPUs (%d) on the host. This will work, but performance will be poor. Are you sure? (yes or no)" %(vcpus, cpu_num))
        try:
            if yes_or_no(res):
                break
            vcpus = int(prompt_for_input("How many VCPUs should be attached?"))
        except ValueError, e:
            print "ERROR: ", e
    if vcpus:
            guest.vcpus = vcpus

def get_keymap(keymap, guest):
    if keymap:
        try:
            guest.keymap = keymap
        except ValueError, e:
            print "ERROR: ", e

def get_disk(disk, size, sparse, guest, hvm, conn):
    # FIXME: need to handle a list of disks at some point
    while 1:
        msg = "What would you like to use as the disk (path)?"
        if not size is None:
            msg = "Please enter the path to the file you would like to use for storage. It will have size %sGB." %(size,)
        disk = prompt_for_input(msg, disk)
        while 1:
            if os.path.exists(disk):
                break
            size = prompt_for_input("How large would you like the disk (%s) to be (in gigabytes)?" %(disk,), size)
            try:
                size = float(size)
                break
            except Exception, e:
                print "ERROR: ", e
                size = None

        try:
            d = virtinst.VirtualDisk(disk, size, sparse = sparse)
            if d.is_conflict_disk(conn) is True:
                while 1:
                    retryFlg = False
                    warnmsg = "Disk %s is already in use by another guest!" % disk
                    res = prompt_for_input(warnmsg + "  Do you really want to use the disk (yes or no)? ")
                    try:
                        if yes_or_no(res) is True:
                            break
                        else:
                            retryFlg = True
                            break
                    except ValueError, e:
                        print "ERROR: ", e
                        continue
                if retryFlg is True:
                    disk = size = None
                    continue
            if d.type == virtinst.VirtualDisk.TYPE_FILE and not(hvm) and virtinst.util.is_blktap_capable():
                d.driver_name = virtinst.VirtualDisk.DRIVER_TAP
        except ValueError, e:
            print "ERROR: ", e
            disk = size = None
            continue

        guest.disks.append(d)
        break

def get_disks(disk, size, sparse, guest, hvm, conn):
    # ensure we have equal length lists 
    if (type(disk) == type(size) == list):
        if len(disk) != len(size):
            print >> sys.stderr, "Need to pass size for each disk"
            sys.exit(1)
    elif type(disk) == list:
        size = [ None ] * len(disk)
    elif type(size) == list:
        disk = [ None ] * len(size)

    if (type(disk) == list):
        map(lambda d, s: get_disk(d, s, sparse, guest, hvm, conn),
            disk, size)
    elif (type(size) == list):
        map(lambda d, s: get_disk(d, s, sparse, guest, hvm, conn),
            disk, size)
    else:
        get_disk(disk, size, sparse, guest, hvm, conn)

def get_network(mac, network, guest):
    if mac == "RANDOM":
        mac = None
    if network == "user":
        n = virtinst.VirtualNetworkInterface(mac, type="user")
    elif network[0:6] == "bridge":
        n = virtinst.VirtualNetworkInterface(mac, type="bridge", bridge=network[7:])
    elif network[0:7] == "network":
        n = virtinst.VirtualNetworkInterface(mac, type="network", network=network[8:])
    else:
        print >> sys.stderr, "Unknown network type " + network
        sys.exit(1)
    guest.nics.append(n)

def get_networks(macs, bridges, networks, guest):
    if type(bridges) != list and bridges != None:
        bridges = [ bridges ]

    if type(macs) != list and macs != None:
        macs = [ macs ]

    if type(networks) != list and networks != None:
        networks = [ networks ]

    if bridges is not None and networks != None:
        print >> sys.stderr, "Cannot mix both --bridge and --network arguments"
        sys.exit(1)

    # ensure we have equal length lists
    if bridges != None:
        networks = map(lambda b: "bridge:" + b, bridges)

    if networks != None:
        if macs != None:
            if len(macs) != len(networks):
                print >> sys.stderr, "Need to pass equal numbers of networks & mac addresses"
                sys.exit(1)
        else:
            macs = [ None ] * len(networks)
    else:
        if os.getuid() == 0:
            net = virtinst.util.default_network()
            networks = [net[0] + ":" + net[1]]
        else:
            networks = ["user"]
        if macs != None:
            if len(macs) > 1:
                print >> sys.stderr, "Need to pass equal numbers of networks & mac addresses"
                sys.exit(1)
        else:
            macs = [ None ]

    map(lambda m, n: get_network(m, n, guest), macs, networks)

def get_graphics(vnc, vncport, nographics, sdl, keymap, guest):
    if vnc and nographics:
        raise ValueError, "Can't do both VNC graphics and nographics"
    elif vnc and sdl:
        raise ValueError, "Can't do both VNC graphics and SDL"
    elif sdl and nographics:
        raise ValueError, "Can't do both SDL and nographics"
    if nographics:
        guest.graphics = False
        return
    if vnc is not None:
        guest.graphics = (True, "vnc", vncport, keymap)
        return
    if sdl is not None:
        guest.graphics = (True, "sdl")
        return
    while 1:
        res = prompt_for_input("Would you like to enable graphics support? (yes or no)")
        try:
            vnc = yes_or_no(res)
        except ValueError, e:
            print "ERROR", e
            continue
        if vnc:
            guest.graphics = (True, "vnc", vncport, keymap)
        else:
            guest.graphics = False
        break


### Paravirt input gathering functions
def get_paravirt_install(src, guest):
    while 1:
        src = prompt_for_input("What is the install location?", src)
        try:
            guest.location = src
            break
        except ValueError, e:
            print "ERROR: ", e
            src = None

def get_paravirt_extraargs(extra, guest):
    guest.extraargs = extra


### fullvirt input gathering functions
def get_fullvirt_cdrom(cdpath, location, guest):
    if cdpath is None and location is not None:
        cdpath = location

    while 1:
        cdpath = prompt_for_input("What is the virtual CD image, CD device or install location?", cdpath)
        try:
            guest.location = cdpath
            break
        except ValueError, e:
            print "ERROR: ", e
            cdpath = None

### Option parsing
def check_before_store(option, opt_str, value, parser):
    if len(value) == 0:
        raise OptionValueError, "%s option requires an argument" %opt_str
    setattr(parser.values, option.dest, value)

def check_before_append(option, opt_str, value, parser):
    if len(value) == 0:
        raise OptionValueError, "%s option requires an argument" %opt_str
    parser.values.ensure_value(option.dest, []).append(value)

def parse_args():
    parser = OptionParser()
    parser.add_option("-n", "--name", type="string", dest="name",
                      action="callback", callback=check_before_store,
                      help="Name of the guest instance")
    parser.add_option("-r", "--ram", type="int", dest="memory",
                      help="Memory to allocate for guest instance in megabytes")
    parser.add_option("-u", "--uuid", type="string", dest="uuid",
                      action="callback", callback=check_before_store,
                      help="UUID for the guest; if none is given a random UUID will be generated. If you specify UUID, you should use a 32-digit hexadecimal number.")
    parser.add_option("", "--vcpus", type="int", dest="vcpus",
                      help="Number of vcpus to configure for your guest")
    parser.add_option("", "--check-cpu", action="store_true", dest="check_cpu",
                      help="Check that vcpus do not exceed physical CPUs and warn if they do.")

    # disk options
    parser.add_option("-f", "--file", type="string",
                      dest="diskfile", action="callback", callback=check_before_append,
                      help="File to use as the disk image")
    parser.add_option("-s", "--file-size", type="float",
                      action="append", dest="disksize",
                      help="Size of the disk image (if it doesn't exist) in gigabytes")
    parser.add_option("", "--nonsparse", action="store_false",
                      default=True, dest="sparse",
                      help="Don't use sparse files for disks.  Note that this will be significantly slower for guest creation")
    
    # network options
    parser.add_option("-m", "--mac", type="string",
                      dest="mac", action="callback", callback=check_before_append,
                      help="Fixed MAC address for the guest; if none or RANDOM is given a random address will be used")
    parser.add_option("-b", "--bridge", type="string",
                      dest="bridge", action="callback", callback=check_before_append,
                      help="Bridge to connect guest NIC to; if none given, will try to determine the default")
    parser.add_option("-w", "--network", type="string",
                      dest="network", action="callback", callback=check_before_append,
                      help="Connect the guest to a virtual network, forwarding to the physical network with NAT")

    # graphics options
    parser.add_option("", "--vnc", action="store_true", dest="vnc", 
                      help="Use VNC for graphics support")
    parser.add_option("", "--vncport", type="int", dest="vncport",
                      help="Port to use for VNC")
    parser.add_option("", "--sdl", action="store_true", dest="sdl", 
                      help="Use SDL for graphics support")
    parser.add_option("", "--nographics", action="store_true",
                      help="Don't set up a graphical console for the guest.")
    parser.add_option("", "--noautoconsole",
                      action="store_false", dest="autoconsole",
                      help="Don't automatically try to connect to the guest console")

    parser.add_option("-k", "--keymap", type="string", dest="keymap",
                      action="callback", callback=check_before_store,
                      help="set up keymap for a graphical console")

    parser.add_option("", "--accelerate", action="store_true", dest="accelerate",
                      help="Use kernel acceleration capabilities")
    parser.add_option("", "--connect", type="string", dest="connect",
                      action="callback", callback=check_before_store,
                      help="Connect to hypervisor with URI",
                      default=virtinst.util.default_connection())
    
    # fullvirt options
    parser.add_option("-v", "--hvm", action="store_true", dest="fullvirt",
                      help="This guest should be a fully virtualized guest")
    parser.add_option("-c", "--cdrom", type="string", dest="cdrom",
                      action="callback", callback=check_before_store,
                      help="File to use a virtual CD-ROM device for fully virtualized guests")
    parser.add_option("", "--os-type", type="string", dest="os_type",
                      action="callback", callback=check_before_store,
                      help="The OS type for fully virtualized guests, e.g. 'linux', 'unix', 'windows'")
    parser.add_option("", "--os-variant", type="string", dest="os_variant",
                      action="callback", callback=check_before_store,
                      help="The OS variant for fully virtualized guests, e.g. 'fedora6', 'rhel5', 'solaris10', 'win2k', 'vista'")
    parser.add_option("", "--noapic", action="store_true", dest="noapic", help="Disables APIC for fully virtualized guest (overrides value in os-type/os-variant db)", default=False)
    parser.add_option("", "--noacpi", action="store_true", dest="noacpi", help="Disables ACPI for fully virtualized guest (overrides value in os-type/os-variant db)", default=False)
    parser.add_option("", "--arch", type="string", dest="arch",
                      action="callback", callback=check_before_store,
                      help="The CPU architecture to simulate")
    
    # paravirt options
    parser.add_option("-p", "--paravirt", action="store_false", dest="paravirt",
                      help="This guest should be a paravirtualized guest")
    parser.add_option("-l", "--location", type="string", dest="location",
                      action="callback", callback=check_before_store,
                      help="Installation source for paravirtualized guest (eg, nfs:host:/path, http://host/path, ftp://host/path)")
    parser.add_option("-x", "--extra-args", type="string",
                      dest="extra", default="",
                      help="Additional arguments to pass to the installer with paravirt guests")

    # Misc options
    parser.add_option("-d", "--debug", action="store_true", dest="debug", 
                      help="Print debugging information")
    parser.add_option("", "--noreboot", action="store_true", dest="noreboot", help="Disables the automatic rebooting when the installation is complete.")

    (options,args) = parser.parse_args()
    return options


### console callback methods
def get_xml_string(dom, path):
    xml = dom.XMLDesc(0)
    try:
        doc = libxml2.parseDoc(xml)
    except:
        return None

    ctx = doc.xpathNewContext()
    try:
        ret = ctx.xpathEval(path)
        tty = None
        if len(ret) == 1:
            tty = ret[0].content
        ctx.xpathFreeContext()
        doc.freeDoc()
        return tty
    except Exception, e:
        ctx.xpathFreeContext()
        doc.freeDoc()
        return None

def vnc_console(dom):
    import time;
    num = 0
    while num < ( 40 / 0.25 ): # 40 seconds, .25 second sleeps
        vncport = get_xml_string(dom,
                                 "/domain/devices/graphics[@type='vnc']/@port")
        if vncport == '-1':
            num += 1
            time.sleep(0.25)
        else:
            break
    if vncport == '-1' or vncport is None:
        print >> sys.stderr, "Unable to connect to graphical console; vnc port number not found."
        return None
    vncport = int(vncport)
    vnchost = "localhost"
    if not os.path.exists("/usr/bin/vncviewer"):
        print >> sys.stderr, "Unable to connect to graphical console; vncviewer not installed.  Please connect to %s:%d" %(vnchost, vncport)
        return None
    if not os.environ.has_key("DISPLAY"):
        print >> sys.stderr, "Unable to connect to graphical console; DISPLAY is not set.  Please connect to %s:%d" %(vnchost, vncport)
        return None
    logging.debug("VNC Port: %d; VNC host: %s" % (vncport, vnchost))

    child = os.fork()
    if not child:
        os.execvp("/usr/bin/vncviewer", ["/usr/bin/vncviewer",
                                         "%s:%d" %(vnchost, vncport) ])
        os._exit(1)

    return child

def txt_console(dom):
    tty = get_xml_string(dom, "/domain/devices/console/@tty")
    if tty is None or not os.access(tty, os.R_OK | os.W_OK):
        return None
    child = os.fork()
    if not child:
        os.execvp("/usr/sbin/xm",
                  ["/usr/sbin/xm", "console", "%s" %(dom.ID(),)])
        os._exit(1)

    return child

def show_console(dom):
    gfxtype = get_xml_string(dom, "/domain/devices/graphics/@type")
    if gfxtype == "vnc":
        return vnc_console(dom)
    return txt_console(dom)

### Let's do it!
def main():
    options = parse_args()

    if options.debug:
        logging.basicConfig(level=logging.DEBUG,
                            format="%(asctime)s %(levelname)-8s %(message)s",
                            datefmt="%a, %d %b %Y %H:%M:%S",
                            stream=sys.stderr)
    else:
        logging.basicConfig(level=logging.ERROR,
                            format="%(asctime)s %(levelname)-8s %(message)s",
                            datefmt="%a, %d %b %Y %H:%M:%S",
                            stream=sys.stderr)

    conn = libvirt.open(options.connect)
    type = None
    # check to ensure we're really on a xen kernel
    if conn.getType() == "Xen":
        type = "xen"

        if os.geteuid() != 0:
            print >> sys.stderr, "Must be root to install guests"
            sys.exit(1)

    # first things first, are we trying to create a fully virt guest?
    if conn.getType() == "Xen":
        hvm = options.fullvirt
        pv  = options.paravirt
        if virtinst.util.is_hvm_capable():
            if hvm is not None and pv is not None:
                print >> sys.stderr, "Can't do both --hvm and --paravirt"
                sys.exit(1)
            elif pv is not None:
                hvm = False
        else:
            if hvm is not None:
                print >> sys.stderr, "Can't do --hvm to this system: HVM guest is not supported by your CPU or enabled in your BIOS"
                sys.exit(1)
            else:
                hvm = False
        if hvm is None:
            hvm = get_full_virt()
    else:
        hvm = True
        type = "qemu"
        if options.accelerate:
            if virtinst.util.is_kvm_capable():
                type = "kvm"
            elif virtinst.util.is_kqemu_capable():
                type = "kqemu"

    if hvm:
        if options.arch is None:
            guest = virtinst.FullVirtGuest(connection=conn, type=type)
        else:
            guest = virtinst.FullVirtGuest(connection=conn, type=type, arch=options.arch)
    else:
        guest = virtinst.ParaVirtGuest(connection=conn, type=type)

    # now let's get some of the common questions out of the way
    get_name(options.name, guest)
    get_memory(options.memory, guest)
    get_uuid(options.uuid, guest)
    get_vcpus(options.vcpus, options.check_cpu, guest, conn)
    get_keymap(options.keymap, guest)

    # set up disks
    get_disks(options.diskfile, options.disksize, options.sparse,
              guest, hvm, conn)

    # set up network information
    get_networks(options.mac, options.bridge, options.network, guest)

    # set up graphics information
    get_graphics(options.vnc, options.vncport, options.nographics, options.sdl, options.keymap, guest)

    # and now for the full-virt vs paravirt specific questions
    if not hvm: # paravirt
        get_paravirt_install(options.location, guest)
        get_paravirt_extraargs(options.extra, guest)
        continue_inst = False
    else:
        get_fullvirt_cdrom(options.cdrom, options.location, guest)
        if options.noacpi:
            guest.features["acpi"] = False
        if options.noapic:
            guest.features["apic"] = False
        if options.os_type is not None:
            guest.set_os_type(options.os_type)
            if options.os_variant is not None:
                guest.set_os_variant(options.os_variant)
        continue_inst = guest.get_continue_inst()
        
    if options.autoconsole is False:
        conscb = None
    else:
        conscb = show_console

    progresscb = progress.TextMeter()

    # we've got everything -- try to start the install
    try:
        print "\n\nStarting install..."
        dom = guest.start_install(conscb,progresscb)
        if dom is None:
            print "Guest installation failed"
            sys.exit(0)
        elif dom.info()[0] != libvirt.VIR_DOMAIN_SHUTOFF:
            # domain seems to be running
            print "Domain installation still in progress.  You can reconnect "
            print "to the console to complete the installation process."
            sys.exit(0)
    except RuntimeError, e:
        print >> sys.stderr, "ERROR: ", e
        sys.exit(1)

    # the domain is no longer running
    # FIXME: this is just a hacky heuristic, but I'll take what I can get
    try:
        fd = os.open(guest.disks[0].path, os.O_RDONLY)
        buf = os.read(fd, 512)
        os.close(fd)
        if len(buf) == 512 and \
               struct.unpack("H", buf[0x1fe: 0x200]) == (0xaa55,):
            # things installed enough that we should be able to restart
            # the domain
            if continue_inst:
                # continue to installation.
                dom = guest.continue_install(conscb,progresscb)
                if dom is None:
                    print "Guest installation failed"
                    sys.exit(0)
                elif dom.info()[0] != libvirt.VIR_DOMAIN_SHUTOFF:
                    # domain seems to be running
                    print "Domain installation still in progress.  You can reconnect "
                    print "to the console to complete the installation process."
                    sys.exit(0)
            if options.noreboot:
                print ("Guest installation complete... you can restart your domain\n"
                      "by running 'virsh start %s'") %(guest.name,)
            else:
                print "Guest installation complete... restarting guest."
                dom.create()
                guest.connect_console(conscb)
        else:
            print ("Domain installation does not appear to have been\n"
                   "successful.  If it was, you can restart your domain\n"
                   "by running 'virsh start %s'; otherwise, please\n"
                   "restart your installation.") %(guest.name,)
    except Exception, e:
        print "exception was:", e
        print ("Domain installation may not have been\n"
               "successful.  If it was, you can restart your domain\n"
               "by running 'virsh start %s'; otherwise, please\n"
               "restart your installation.") %(guest.name,)


if __name__ == "__main__":
    main()
