#!/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 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.XenDisk(disk, size, sparse = sparse)
            if d.type == virtinst.XenDisk.TYPE_FILE and not(hvm) and virtinst.util.is_blktap_capable():
                d.driver_name = virtinst.XenDisk.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.XenNetworkInterface(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.cdrom = 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")
    
    
    # vmx/svm options
    if virtinst.util.is_hvm_capable():
        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")

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

    # check to ensure we're really on a xen kernel
    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 virtinst.util.is_hvm_capable():
        hvm = options.fullvirt
    if hvm is None:
        hvm = get_full_virt()
    if hvm:
        guest = virtinst.FullVirtGuest()
    else:
        guest = virtinst.ParaVirtGuest()

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

    if options.autoconsole is False:
        conscb = None
    else:
        conscb = show_console

    # we've got everything -- try to start the install
    try:
        print "\n\nStarting install..."
        dom = guest.start_install(conscb)
        if dom is not None and guest.domain.info()[0] != 0:
            # 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."
            guest.start_from_disk(conscb)
        else:
            print ("Domain installation does not appear to have been\n"
                   "successful.  If it was, you can restart your domain\n"
                   "by running 'xm create -c %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 'xm create -c %s'; otherwise, please\n"
               "restart your installation.") %(guest.name,)


if __name__ == "__main__":
    main()
