#!/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
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()


def check_xen():
    if not os.path.isdir("/proc/xen"):
        print >> sys.stderr, "Can only install guests if running under a Xen kernel!"
        sys.exit(1)

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

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

def get_disk(disk, size, sparse, guest, hvm):
    # FIXME: need to handle a list of disks at some point
    while 1:
        disk = prompt_for_input("What would you like to use as the disk (path)?", 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.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):
    # 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)

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

def get_network(mac, bridge, guest):
    if mac == "RANDOM":
        mac = None
    n = virtinst.VirtualNetworkInterface(mac, bridge)
    guest.nics.append(n)

def get_networks(macs, bridges, guest):
    # ensure we have equal length lists 
    if (type(macs) == type(bridges) == list):
        if len(macs) != len(bridges):
            print >> sys.stderr, "Need to pass bridge for each network device"
            sys.exit(1)
    elif type(macs) == list:
        bridges = [ None ] * len(macs)
    elif type(bridges) == list:
        macs = [ None ] * len(bridges)

    if (type(macs) == list):
        map(lambda m, b: get_network(m, b, guest), macs, bridges)
    else:
        get_network(macs, bridges, guest)

def get_graphics(vnc, vncport, nographics, sdl, guest):
    if vnc and nographics:
        raise ValueError, "Can't do both VNC graphics and nographics"
    if nographics:
        guest.graphics = False
        return
    if vnc is not None:
        guest.graphics = (True, "vnc", vncport)
        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 = "vnc"
        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, guest):
    while 1:
        cdpath = prompt_for_input("What would you like to use for the virtual CD image?", cdpath)
        try:
            guest.location = cdpath
            break
        except ValueError, e:
            print "ERROR: ", e
            cdpath = None

### Option parsing
def parse_args():
    parser = OptionParser()
    parser.add_option("-n", "--name", type="string", dest="name",
                      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",
                      help="UUID for the guest; if none is given a random UUID will be generated")
    parser.add_option("", "--vcpus", type="int", dest="vcpus",
                      help="Number of vcpus to configure for your guest")

    # disk options
    parser.add_option("-f", "--file", type="string",
                      dest="diskfile", action="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="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="append",
                      help="Bridge to connect guest NIC to; if none given, will try to determine the default")

    # 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("", "--accelerate", action="store_true", dest="accelerate",
                      help="Use kernel acceleration capabilities")
    parser.add_option("", "--connect", type="string", dest="connect",
                      help="Connect to hypervisor with URI")
    
    # 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",
                      help="File to use a virtual CD-ROM device for fully virtualized guests")
    parser.add_option("", "--os-type", type="string", dest="os_type", help="The OS type for fully virtualized guests, e.g. Linux, Solaris, Windows", default="Other")
    parser.add_option("", "--os-variant", type="string", dest="os_variant", help="The OS variant for fully virtualized guests, e.g. Fedora, Solaris 8, Windows XP", default="Other")
    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", help="The CPU architecture to simulate")
    
    # paravirt options
    parser.add_option("-p", "--paravirt", action="store_false", dest="fullvirt",
                      help="This guest should be a paravirtualized guest")
    parser.add_option("-l", "--location", type="string", dest="location",
                      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")


    (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; time.sleep(2) # FIXME: ugh. 
    vncport = get_xml_string(dom,
                             "/domain/devices/graphics[@type='vnc']/@port")
    if vncport == None:
        vncport = 5900 + dom.ID()
    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

    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"
        check_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?
    hvm = False
    if conn.getType() == "xen":
        if virtinst.util.is_hvm_capable():
            hvm = options.fullvirt
        if hvm is None:
            hvm = get_full_virt()
    else:
        hvm = True
        type = "qemu"
        if options.accelerate:
            if util.is_kvm_capable():
                type = "kvm"
            elif 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, guest)

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

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

    # set up graphics information
    get_graphics(options.vnc, options.vncport, options.nographics, options.sdl, 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)
    else:
        get_fullvirt_cdrom(options.cdrom, guest)
        guest.set_os_type(options.os_type)
        guest.set_os_variant(options.os_variant)
        if options.noacpi:
            guest.features["acpi"] = False
        if options.noapic:
            guest.features["apic"] = False
        
    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
            print "Guest installation complete... restarting guest."
            dom.start()
        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()
