#!/usr/bin/env python
##
# Bootloader for EL-based distros that support Xen. 
#
# We keep all logic for booting these distros contained within this file
# where possible.

# Some brief documentation of the other-config keys this tool knows about:
#
# install-repository: Required.  Path to a repository; 'http', 'https', or 
#    'nfs'.  Should be specified as would be used by the target installer, not
#    including prefixes, e.g. method=.
#
# install-vnc:  Default: false.  Use VNC where available during the
#    installation.  
#
# install-vncpasswd:  Default: empty.  The VNC password to use, when providing
#    one is possible via the command-line of the target distro
#
# install-round:  Default: 1.  The current bootloader round.  Not to be edited
#    by the user
#
# install-distro:  Default: 'rhlike'.  The general distribution type.  Currently
#    supported values are 'rhlike', 'sleslike', and 'debianlike'.
#
# install-arch:  Default: i386.  The architecture to install.

import sys
import subprocess
import os
import shutil
import getopt
import xmlrpclib
import httplib
import socket
import tempfile
import urllib2
import gzip
import traceback
import syslog
import md5

sys.path.append("/usr/lib/python")

BOOTDIR = "/var/run/xend/boot"
PYGRUB = "/usr/bin/pygrub"

RPC_SUCCESS = "Success"

(
    DISTRO_RHLIKE,
    DISTRO_SLESLIKE,
    DISTRO_DEBIANLIKE
) = range(3)

distros = { "rhlike" : DISTRO_RHLIKE, "sleslike" : DISTRO_SLESLIKE, "debianlike": DISTRO_DEBIANLIKE }

rounds = {
    DISTRO_RHLIKE: 1, DISTRO_SLESLIKE: 2, DISTRO_DEBIANLIKE: 1
    }

guest_installer_dir = "/opt/xensource/packages/files/guest-installer"

# We can sometimes tweak an installer's initrd to give it extra features, e.g.
# CD installs in PV guests.  This dictionary specifies the cpio archive used
# as a source for adding files from given an initrd whose md5sum when taken
# directly from the vendor's CD as the key.  Note that we use cpio archives
# because although it doesn't seem to be supported yet, it might that in future
# we can simply cat multiple archives together into a single image, and the
# Linux loader will do the right thing.

# Later initrds are cpio.gz archives
cpio_initrd_fixups = {
    "20d2e6ee93ec377fe3853766751cd94c": "el5.3-initrd-additions.cpio", # RHEL   5.3 x86_64
    "07ab25b6779096b131d97ba168443cb3": "el5.3-initrd-additions.cpio", # RHEL   5.3
    "682896abf6edb88447c66f7734b9edb4": "el5.2-initrd-additions.cpio", # RHEL   5.2 x86_64
    "a7af2795e37e525547aa4262ff462469": "el5.2-initrd-additions.cpio", # RHEL   5.2
    "5fc88d3847596477b522888564a69dff": "el5.1-initrd-additions.cpio", # RHEL   5.1 x86_64
    "8d1e325a2d3c6141fe41d3a214af0993": "el5.1-initrd-additions.cpio", # RHEL   5.1
    "57fd28552875789a3aed79d10afee320": "el5.0-initrd-additions.cpio", # RHEL   5   x86_64
    "367bcaf30b0a411d6b1519947b927542": "el5.0-initrd-additions.cpio", # RHEL   5
    "f1dfaa1d8238211932bf272976348638": "el5.2-initrd-additions.cpio", # Oracle 5.2 x86_64
    "60a31c3c937a1d19575cca86f5330183": "el5.2-initrd-additions.cpio", # Oracle 5.2
    "a1eb076f57d1495e0234038ab29db84d": "el5.1-initrd-additions.cpio", # Oracle 5.1 x86_64
    "c13f4bdfa3a0d90d285571467068d324": "el5.1-initrd-additions.cpio", # Oracle 5.1
    "e1f60fafa660c4537f432377d166f3a2": "el5.0-initrd-additions.cpio", # Oracle 5   x86_64
    "906a03a25a8c003f110f7ebfdf30fae1": "el5.0-initrd-additions.cpio", # Oracle 5
    "62634c97165423d0e76a94962382e5fd": "el5.3-initrd-additions.cpio", # CentOS 5.3 x86_64
    "ec53dc76b4cb77a5331dd2144a49ff7f": "el5.3-initrd-additions.cpio", # CentOS 5.3
    "6b5b1388cd65e51b142b50cde711b9cc": "el5.2-initrd-additions.cpio", # CentOS 5.2 x86_64
    "798baa3b76b295e4f13348f11238bb04": "el5.2-initrd-additions.cpio", # CentOS 5.2
    "b70f6cbf5f9043c83baba7c41ef3bfa2": "el5.1-initrd-additions.cpio", # CentOS 5.1 x86_64
    "118c1f3142ee8abca4ef0d2075ab88a9": "el5.1-initrd-additions.cpio", # CentOS 5.1
    "a14dd1943ba6c4100351713cf55130af": "el5.0-initrd-additions.cpio", # CentOS 5.0 x86_64
    "4610250a8d1565ba4db28be710470b16": "el5.0-initrd-additions.cpio", # CentOS 5.0
}
# Earlier initrds are ext2.gz filesystems
ext2_initrd_fixups = {
    "d9c6c73beedc1eba8710c3af014125f1": "el4.7-initrd-additions.cpio", # RHEL 4.7
    "e56b74ea16503ae28edd9ed3966ca582": "el4.6-initrd-additions.cpio", # RHEL 4.6
    "86358ff197e3b9761e758f6ad0cc21b3": "el4.5-initrd-additions.cpio", # RHEL 4.5
    "bce856015dd4b1589f9bd608398fd036": "el4.7-initrd-additions.cpio", # Cent0S 4.7
    "19d4c41873d28a937cfc2bf381cbcc68": "el4.6-initrd-additions.cpio", # Cent0S 4.6
    "6d192a17a3e2fc280747a948b0b91999": "el4.5-initrd-additions.cpio", # CentOS 4.5
}

##### begin hack.  Provide xmlrpc over UNIX domain socket:
class UDSHTTPConnection(httplib.HTTPConnection):
    """ Stupid hacked up HTTPConnection subclass to allow HTTP over Unix domain
    sockets. """
    def connect(self):
        path = self.host.replace("_", "/")
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(path)

class UDSHTTP(httplib.HTTP):
    _connection_class = UDSHTTPConnection

class UDSTransport(xmlrpclib.Transport):
    def make_connection(self, host):
        return UDSHTTP(host)

def xapi_local():
    return xmlrpclib.Server("http://_var_xapi_xapi/", transport=UDSTransport())
##### end hack.


#### EXCEPTIONS

class UsageError(Exception):
    pass

class APILevelException(Exception):
    exname = "INTERNAL_ERROR"
    def apifmt(self):
        rc = self.exname + "\n"
        for a in self.args:
            rc += a + "\n"
        return rc

class UnsupportedInstallMethod(APILevelException):
    exname = "UNSUPPORTED_INSTALL_METHOD"

class SupportPackageMissing(APILevelException):
    exname = "SUPPORT_PACKAGE_MISSING"

class UnknownApiError(APILevelException):
    exname = "UNKNOWN_API_ERROR"

class InvalidSource(APILevelException):
    exname = "INVALID_SOURCE"

class ResourceNotFound(Exception):
    def __init__(self, name):
        self.name = name

class MountFailureException(Exception):
    pass

##### UTILITY FUNCTIONS

def mount(dev, mountpoint, options = None, fstype = None):
    cmd = ['/bin/mount']
    if options:
        assert type(options) == list

    if fstype:
        cmd.append('-t')
        cmd.append(fstype)

    if options:
        cmd.append("-o")
        cmd.append(",".join(options))

    cmd.append(dev)
    cmd.append(mountpoint)

    rc = subprocess.Popen(cmd, stdout = subprocess.PIPE,
                          stderr = subprocess.PIPE).wait()
    if rc != 0:
        raise MountFailureException, cmd

def umount(mountpoint, force = False):
    rc = subprocess.Popen(['/bin/umount', mountpoint],
                          stdout = subprocess.PIPE,
                          stderr = subprocess.PIPE).wait()
    return rc

# Creation of an NfsRepo object triggers a mount, and the mountpoint is stored int obj.mntpoint.
# The umount is done automatically when the object goes out of scope
class NfsRepo:
    # repo is nfs:server:/path/to/repo or nfs://server/path/to/repo or nfs://server:/path/to/repo
    def __init__(self, repo):
        self.mntpoint = None

        # we deal with RHEL-like NFS paths - if it's a SLES one then
        # turn it into something we can understand first:
        if repo.startswith("nfs://"):
            rest = repo[6:]
            if not "/" in rest:
                raise InvalidSource, "NFS path was not in a valid format"
            server, path = rest.split("/", 1)
            server = server.rstrip(":")
            repo = "nfs:%s:/%s" % (server, path)

        # work out the components:
        [_, server, dir] = repo.split(':')
        if dir[0] != '/':
            raise InvalidSource, "Directory part of NFS path was not an absolute path."

        # make a mountpoint:
        self.mntpoint = tempfile.mkdtemp(dir = '/tmp', prefix = 'nfs-repo-')
        try:
            mount('%s:%s' % (server, dir), self.mntpoint, fstype = "nfs", options = ['ro'])
        except MountFailureException, e:
            # Mount failed.  Re-raise as InvalidSource.
            umount(self.mntpoint)
            os.rmdir(self.mntpoint)
            self.mntpoint = None
            raise InvalidSource, "nfs repo %s" % repo
    
    def __del__(self):
        # if we're getting called due to an unhandled exception, the 
        # os module may have already been unloaded
        import os
        if self.mntpoint:
            os.system("umount %s" % self.mntpoint)
            os.rmdir(self.mntpoint)
    
# Creation of an CdromRepo object triggers a mount, and the mountpoint is stored int obj.mntpoint.
# The umount is done automatically when the object goes out of scope
class CdromRepo:
    # img is a dev node
    def __init__(self, img):
        # make a mountpoint:
        self.mntpoint = tempfile.mkdtemp(dir = '/tmp', prefix = 'cdrom-repo-')
        try:
            mount(img, self.mntpoint, fstype = "iso9660", options = ['ro'])
        except MountFailureException, e:
            # Mount failed.  Re-raise as InvalidSource.
            umount(self.mntpoint)
            os.rmdir(self.mntpoint)
            self.mntpoint = None
            raise InvalidSource, "cdrom repo %s" % img
    
    def __del__(self):
        # if we're getting called due to an unhandled exception, the 
        # os module may have already been unloaded
        import os
        if self.mntpoint:
            os.system("umount %s" % self.mntpoint)
            os.rmdir(self.mntpoint)

# Modified from host-installer.hg/util.py to support not_really flag
# source may be
#  http://blah
#  ftp://blah
#  file://blah
#
# Raises ResourceNotFound or InvalidSource.
#
# If (not_really == True) do not actually copy the file, 
# just return True for "exists" or False for "does not exist"
#
def fetchFile(source, dest = None, not_really = False):

    if source[:5] == 'http:' or source[:5] == 'file:' or source[:4] == 'ftp:':
        # This something that can be fetched using urllib2:

        if not_really:
            # We are just testing for existence
            try:
                fd = urllib2.urlopen(source)
                fd.close()
            except:
                return False
            else:
                return True
            
        # Actually get the file
        try:
            fd = urllib2.urlopen(source)
        except OSError, e:
            # file not found? (from file://)
            if e.errno == 2:
                raise ResourceNotFound, source
            else:
                # something else, we'll re-raise:
                raise
        except urllib2.HTTPError, e:
            # file not found?
            if e.code == 404:
                raise ResourceNotFound, source
            else:
                # something else, we'll re-raise:
                raise
        except urllib2.URLError, e:
            # bad hostname, malformed URL, etc.
            raise ResourceNotFound, source
        except IOError, e:
            # file not found? (from ftp://)
            if e.errno == "ftp error":
                raise ResourceNotFound, source
            else:
                raise
        fd_dest = open(dest, 'w')
        shutil.copyfileobj(fd, fd_dest)
        fd_dest.close()
        fd.close()
    else:
        raise InvalidSource, "Unknown source type."

def close_mkstemp(dir = None, prefix = 'tmp'):
    fd, name = tempfile.mkstemp(dir = dir, prefix = prefix)
    os.close(fd)
    return name

def canonicaliseOtherConfig(vm_uuid):
    server = xapi_local()
    rc = server.session.login_with_password("root", "")
    if rc['Status'] !=  RPC_SUCCESS:
        raise RuntimeError, "Unable to login to API server"
    session = rc['Value']

    rc = server.VM.get_by_uuid(session, vm_uuid)
    if rc['Status'] !=  RPC_SUCCESS:
        raise RuntimeError, "Unable to get VM records for UUID %s" % vm_uuid
    vm = rc['Value']

    rc = server.VM.get_other_config(session, vm)
    if rc['Status'] !=  RPC_SUCCESS:
        raise RuntimeError, "Unable to get other config map from ref %s, which should be VM UUID %s" % (vm, vm_uuid)
    other_config = rc['Value']

    def collect(d, k, default = None):
        if d.has_key(k):
            return d[k]
        else:
            return default
    rc = { 'install-repository': collect(other_config, 'install-repository'),
           'install-vnc':        collect(other_config, 'install-vnc', "false") in ["1", "true"],
           'install-vncpasswd':  collect(other_config, 'install-vncpasswd'),
           'install-distro':     collect(other_config, 'install-distro', 'rhlike'), 
           'install-round':      collect(other_config, 'install-round', '1'),
           'install-arch':       collect(other_config, 'install-arch', 'i386'),
           'debian-release':     collect(other_config, 'debian-release') }
    return rc

def switchBootloader(vm_uuid, target_bootloader = "pygrub"):
    server = xapi_local()
    rc = server.session.login_with_password("root", "")
    if rc['Status'] !=  RPC_SUCCESS:
        raise RuntimeError, "Unable to login to API server"
    session = rc['Value']

    rc = server.VM.get_by_uuid(session, vm_uuid)
    if rc['Status'] !=  RPC_SUCCESS:
        raise RuntimeError, "Unable to get VM records for UUID %s" % vm_uuid
    vm = rc['Value']

    rc = server.VM.set_PV_bootloader(session, vm, target_bootloader)
    if rc['Status'] != RPC_SUCCESS:
        raise RuntimeError, "Unable to change bootloader."

    server.session.logout(session)

def unpack_cpio_initrd(filename, working_dir, check_compressed = True):
    # we'll assume it's a gzipped cpio for now...
    cpio_archive = close_mkstemp(dir = "/tmp", prefix = "initrd-")
    if check_compressed:
        # XXX assume it's compressed for now:
        gz = gzip.GzipFile(filename)
    else:
        gz = open(filename)
    cpio = subprocess.Popen(["/bin/cpio", "-idu", "--quiet"], cwd = working_dir,
                            stdin = subprocess.PIPE)
    while True:
        data = gz.read(1024 * 256)
        if data == "":
            break
        cpio.stdin.write(data)
    cpio.communicate()
    gz.close()

def mount_ext2_initrd(infile, outfile, working_dir):
    # we'll assume it's a gzipped ext2 f/s for now...
    fd = open(outfile, "w")
    gz = gzip.GzipFile(infile)
        
    while True:
        data = gz.read(1024 * 256)
        if data == "":
            break
        fd.write(data)

    fd.close()
    gz.close()

    mount(outfile, working_dir, options = ['loop'])

def md5sum(filename):
    """ Compute the md5sum of a file.  string -> string. """
    fd = open(filename, "r")
    try:
        sumobj = md5.new()
        data = ""
        while True:
            data = fd.read(1024 * 1024)
            if data == "":
                break
            sumobj.update(data)
        return sumobj.hexdigest()
    finally:
        fd.close()

#### INITRD TWEAKING

def mkcpio(working_dir, output_file):
    """ Make a cpio archive containg the files in working_dir, writing the 
    archive to output_file.  It will be uncompressed. """

    # set output_file to be a full path so that we don't create the output
    # file under the new working directory of the cpio process.
    output_file = os.path.realpath(output_file)
    cpio = subprocess.Popen(["/bin/cpio", "-F", output_file, "-oH", "newc", 
                             "--quiet"], cwd = working_dir,
                            stdin = subprocess.PIPE, stdout = None)

    for root, ds, files in os.walk(working_dir):
        assert root.startswith(working_dir), "Root of current walk path starts with original walk path"
	base = root[len(working_dir) + 1:]
        for f in files + ds:
            path = os.path.join(base, f)
            cpio.stdin.write(path + "\n")

    cpio.communicate()

def tweak_initrd(filename):
    """ Patch an initrd with custom files if they are available.  Returns the
    filename of a patched initrd that should be used instead of the file as
    passed in as filename.  The caller is responsible for removing the old
    version of the initrd. """

    digest = md5sum(filename)
    initrd_path = None
    
    if cpio_initrd_fixups.has_key(digest):
        # we can patch this initrd, let's unpack it to a temporary directory:
        working_dir = tempfile.mkdtemp(dir = "/tmp", prefix = "initrd-fixup-")
        cpio_overlay = os.path.join(guest_installer_dir, cpio_initrd_fixups[digest])
        if not os.path.isfile(cpio_overlay):
            raise SupportPackageMissing, "Dom0 does not contain a required file: %s" % cpio_overlay

        # unpack the vendor initrd, then unpack our changes over it:
        unpack_cpio_initrd(filename, working_dir, True)
        unpack_cpio_initrd(cpio_overlay, working_dir, False)

        # now repack to make the final image:
        initrd_path = close_mkstemp(dir = BOOTDIR, prefix="tweaked-initrd-")
        mkcpio(working_dir, initrd_path)

        # clean up the working_dir tree
        shutil.rmtree(working_dir)

    elif ext2_initrd_fixups.has_key(digest):
        # we can patch this initrd, let's unpack it to a temporary directory:
        working_dir = tempfile.mkdtemp(dir = "/tmp", prefix = "initrd-fixup-")
        cpio_overlay = os.path.join(guest_installer_dir, ext2_initrd_fixups[digest])
        if not os.path.isfile(cpio_overlay):
            raise SupportPackageMissing, "Dom0 does not contain a required file: %s" % cpio_overlay
        
        # unpack the vendor initrd, then unpack our changes over it:
        initrd_path = close_mkstemp(dir = BOOTDIR, prefix="tweaked-initrd-")
        mount_ext2_initrd(filename, initrd_path, working_dir)
        unpack_cpio_initrd(cpio_overlay, working_dir, False)
        umount(working_dir)

    return initrd_path

def tweak_bootable_disk(vm):
    xapi = xapi_local()
    rc = xapi.session.login_with_password("root", "")
    if rc['Status'] != RPC_SUCCESS:
        raise UnknownApiError, "Couldn't log in"
    session = rc['Value']

    # get all VBDs, set bootable = (device == 0):
    rc = xapi.VM.get_by_uuid(session, vm)
    if rc['Status'] != RPC_SUCCESS:
        raise UnknownApiError, "Couldn't get VM record"
    vm_ref = rc['Value']

    rc = xapi.VM.get_VBDs(session, vm_ref)
    if rc['Status'] != RPC_SUCCESS:
        raise UnknownApiError, "Couldn't get VBDs"
    
    vbds = rc['Value']
    
    for vbd in vbds:
        rc = xapi.VBD.get_userdevice(session, vbd)
        if rc['Status'] != RPC_SUCCESS:
            raise UnknownApiError, "get_userdevice failed"
        userdevice = rc['Value']
        xapi.VBD.set_bootable(session, vbd, userdevice == "0")

    xapi.session.logout(session)

##### DISTRO-SPECIFIC CODE

def rhel_first_boot_handler(vm, repo_url, 
                        vmlinuz_suburl = "images/xen/vmlinuz",
                        ramdisk_suburl = "images/xen/initrd.img"):

    vmlinuz_file = close_mkstemp(dir = BOOTDIR, prefix = "vmlinuz-")
    ramdisk_file = close_mkstemp(dir = BOOTDIR, prefix = "ramdisk-")

    # download the kernel and ramdisk:
    vmlinuz_url = repo_url + vmlinuz_suburl
    ramdisk_url = repo_url + ramdisk_suburl
    try:
        fetchFile(vmlinuz_url, vmlinuz_file)
        fetchFile(ramdisk_url, ramdisk_file)
    except ResourceNotFound, e:
        os.unlink(vmlinuz_file)
        os.unlink(ramdisk_file)
        raise InvalidSource, "Unable to access a required file in the specified repository: %s." % e.name

    # Possibly apply tweaks to initrd.
    #
    # Currently, this adds support for graphical installation via XenCenter, and fixes
    # installation via ISO (by making loader & anaconda recognise r/o blockdevs on xenbus 
    # as CDROM drives).  However, it could be used for more in future.
    #
    modified_ramdisk = tweak_initrd(ramdisk_file)
    if modified_ramdisk:
        os.unlink(ramdisk_file)
        ramdisk_file = modified_ramdisk

    return vmlinuz_file, ramdisk_file

# Return the extra arg needed by RHEL kernel for it to locate installation media
def rhel_first_boot_args(repo):
    if True in [ repo.startswith(x) for x in ("http", "ftp", "nfs") ]:
        if not repo.endswith("/"):
            return "method=%s/" % repo
    return "method=%s" % repo

def sles_first_boot_handler(vm, repo_url, other_config):

    # look for the xen kernel and initrd in boot first:
    if other_config['install-arch'] == 'x86_64':
        bootdir =      'boot/x86_64/'
        kernel_fname = 'vmlinuz-xen'
        initrd_fname = 'initrd-xen'
    else:
        bootdir =      'boot/i386/'
        kernel_fname = 'vmlinuz-xenpae'
        initrd_fname = 'initrd-xenpae'
        
    vmlinuz_url = repo_url + bootdir + kernel_fname
    vmlinuz_file = close_mkstemp(dir = BOOTDIR, prefix = "vmlinuz-")
    ramdisk_url = repo_url + bootdir + initrd_fname
    ramdisk_file = close_mkstemp(dir = BOOTDIR, prefix = "ramdisk-")
    try:
        fetchFile(vmlinuz_url, vmlinuz_file)
        fetchFile(ramdisk_url, ramdisk_file)
    except ResourceNotFound, e:
        os.unlink(vmlinuz_file)
        os.unlink(ramdisk_file)
        raise InvalidSource, "The repository specified did not contain a required file, %s." % e.name
    
    print >> sys.stderr, "kernel %s and initrd %s from %s used" % (kernel_fname, initrd_fname, bootdir)
    
    return vmlinuz_file, ramdisk_file

# Return the extra arg needed by SLES kernel for it to locate installation media
def sles_first_boot_args(repo):
    if repo == "cdrom":
        # this should really be "install=cd", but the sles installer interprets
        # the disk as a harddrive.  If/when we fix that, we will need to replace
        # the following line.
        return "install=hd"
    if True in [ repo.startswith(x) for x in ("http", "ftp", "nfs") ]:
        if not repo.endswith("/"):
            return "install=%s/" % repo
    return "install=%s" % repo

def debian_first_boot_handler(vm, repo_url, other_config):

    if not other_config.has_key('debian-release'):
        raise UnsupportedInstallMethod, \
            "other-config:debian-release was not set to an appropriate value, " \
            "and this is required for the selected distribution type."
    if not other_config.has_key('install-arch'):
        raise UnsupportedInstallMethod, \
            "other-config:install-arch was not set to an appropriate value, " \
            "and this is required for the selected distribution type."

    if other_config['install-repository'] == "cdrom":
        cdrom_dirs = { 'i386': 'install.386/xen/',
	               'amd64': 'install.amd/xen/' }
        vmlinuz_url = repo_url + cdrom_dirs[other_config['install-arch']] + "vmlinuz"
        ramdisk_url = repo_url + cdrom_dirs[other_config['install-arch']] + "initrd.gz"
    else:
        boot_dir = "dists/%s/main/installer-%s/current/images/netboot/xen/" % (other_config['debian-release'], other_config['install-arch'])
        vmlinuz_url = repo_url + boot_dir + "vmlinuz"
        ramdisk_url = repo_url + boot_dir + "initrd.gz"

    # download the kernel and ramdisk:
    vmlinuz_file = close_mkstemp(dir = BOOTDIR, prefix = "vmlinuz-")
    ramdisk_file = close_mkstemp(dir = BOOTDIR, prefix = "ramdisk-")

    try:
        fetchFile(vmlinuz_url, vmlinuz_file)
        fetchFile(ramdisk_url, ramdisk_file)
    except ResourceNotFound, e:
        os.unlink(vmlinuz_file)
        os.unlink(ramdisk_file)
        raise InvalidSource, "Unable to access a required file in the specified repository: %s." % e.name

    # Possibly apply tweaks to initrd.
    modified_ramdisk = tweak_initrd(ramdisk_file)
    if modified_ramdisk:
        os.unlink(ramdisk_file)
        ramdisk_file = modified_ramdisk

    return vmlinuz_file, ramdisk_file

def debian_first_boot_args(repo):
    return ""

##### MAIN HANDLERS

def handle_first_boot(vm, img, args, other_config):
    if other_config['install-distro'] not in distros.keys():
        raise RuntimeError, "other-config:install-distro was not present or known."
    distro = distros[other_config['install-distro']]

    repo = other_config['install-repository']
    vnc = other_config['install-vnc']
    vncpasswd = other_config['install-vncpasswd']

    # extract the kernel and ramdisk

    # sanity check repo:
    if repo == "cdrom": 
        pass
    elif repo and True in [repo.startswith(x) for x in ['http', 'ftp', 'nfs']]: 
        pass
    else:
        raise UnsupportedInstallMethod, \
            "other-config:install-repository was not set to an appropriate value, " \
            "and this is required for the selected distribution type."

    # calculate repo_url, a prefix that can be passed into fetchFile
    if repo == "cdrom":
        # CdromRepo.__init__ triggers a mount.  CdromRepo.__del__ does the umount.
        cdrom_repo = CdromRepo(img)
        repo_url = "file://%s/" % cdrom_repo.mntpoint
    elif repo.startswith("nfs"):
        # NfsRepo.__init__ triggers a mount.  NfsRepo.__del__ does the umount.
        nfs_repo = NfsRepo(repo)
        repo_url = "file://%s/" % nfs_repo.mntpoint
    else:
        repo_url = repo
        if not repo_url.endswith("/"):
            repo_url += "/"

    # invoke distro specific handler for extraction of kernel and ramdisk
    if distro == DISTRO_RHLIKE:
        kernel, ramdisk = rhel_first_boot_handler(vm, repo_url)
    elif distro == DISTRO_SLESLIKE:
        kernel, ramdisk = sles_first_boot_handler(vm, repo_url, other_config)
    elif distro == DISTRO_DEBIANLIKE:
        kernel, ramdisk = debian_first_boot_handler(vm, repo_url, other_config)
    else:
        raise UnsupportedInstallMethod

    if repo == 'cdrom':
        # SLES/RHEL: booting from CDROM this time but booting from 1st disk next time
        tweak_bootable_disk(vm)

    # Calculate the extra args need by kernel to locate installation repository
    if distro == DISTRO_RHLIKE:
        extra_args = rhel_first_boot_args(repo)
    elif distro == DISTRO_SLESLIKE:
        extra_args = sles_first_boot_args(repo)
    elif distro == DISTRO_DEBIANLIKE:
        extra_args = debian_first_boot_args(repo)
    else:
        raise UnsupportedInstallMethod

    # Tell eliloader to run 2nd boot phase next time this vm is started
    if rounds[distro] == 1:
        switchBootloader(vm)

    # Put it all together
    args += " " + extra_args
    if vnc:
        args += " vnc"
    if vncpasswd:
        args += " vncpassword=%s" % vncpasswd
    print 'linux (kernel %s)(ramdisk %s)(args "%s")' % (kernel, ramdisk, args)

def handle_second_boot(vm, img, args, other_config):
    distro = distros[other_config['install-distro']]    

    prepend_args = [PYGRUB]

    if distro == DISTRO_SLESLIKE:

      # SLES 9/10 installers do not create /boot/grub/menu.lst when installing on top of XEN
      # SLES 11 does not have this problem.
      # If pygrub with no options fails then this must be one of the problematic versions, in
      # which case /we/ need to tell pygrub where to find the kernel and initrd.
      cmd = "pygrub -n %s &>/dev/null" % img
      rc = os.system(cmd)
      if rc != 0:

        # need to emulate domUloader.  This is done by finding a kernel that
        # we can boot if possible, and then setting PV-bootloader-args.
        if other_config['install-arch'] == 'x86_64':
            kernel = 'vmlinuz-xen'
            initrd = 'initrd-xen'
        else:
            kernel = 'vmlinuz-xenpae'
            initrd = 'initrd-xenpae'

        for k, i in [ ("/%s" % kernel, "/%s" % initrd ), ("/boot/%s" % kernel , "/boot/%s" % initrd ) ]:
            print >>sys.stderr, "SLES 10 WORKAROUND: Trying %s and %s" % (k, i)
            cmd = "pygrub -n --kernel %s --ramdisk %s %s &>/dev/null" % (k, i, img)
            print >>sys.stderr, "Trying %s" % cmd
            rc = os.system(cmd)
            if rc == 0:
                # found it - make the setting permanent:
                xapi = xapi_local()
                rc = xapi.session.login_with_password("root", "")
                if rc['Status'] != RPC_SUCCESS:
                    raise UnknownApiError
                session = rc['Value']
                try:
                    rc = xapi.VM.get_by_uuid(session, vm)
                    if rc['Status'] != RPC_SUCCESS:
                        raise UnknownApiError, "Couldn't get VM reference for %s" % vm
                    vm_ref = rc['Value']
                
                    xapi.VM.set_PV_bootloader_args(session, vm_ref, "--kernel %s --ramdisk %s" % (k, i))

                    prepend_args += ["--kernel", k, "--ramdisk", i]
                finally:
                    xapi.session.logout(session)
    else:
        raise UnsupportedInstallMethod


    # now exec pygrub - hackily call update_rounds since we won't get to
    # run again.
    switchBootloader(vm)
    update_rounds(vm, 2, 2)
    os.execv(PYGRUB, prepend_args + sys.argv[1:])

def update_rounds(vm, current_round, rounds_required):
    xapi = xapi_local()
    rc = xapi.session.login_with_password("root", "")
    if rc['Status'] != RPC_SUCCESS:
        raise UnknownApiError, "Couldn't log in"
    session = rc['Value']

    try:
        rc = xapi.VM.get_by_uuid(session, vm)
        if rc['Status'] != RPC_SUCCESS:
            raise UnknownApiError, "Couldn't get VM reference for %s" % vm
        vm_ref = rc['Value']

        # remove the install-round field: ignore errors as the key might
        # not be there and this is OK (default value is 1).
        xapi.VM.remove_from_other_config(session, vm_ref, "install-round")
    
        # write a new value in for install-round if appropriate:
        if current_round != rounds_required:
            rc = xapi.VM.add_to_other_config(session, vm_ref, "install-round", str(current_round + 1))
            if rc['Status'] != RPC_SUCCESS:
                raise UnknownApiError, "Couldn't set other-config:install-round."
        else:
            # All rounds complete. Remove install-distro key from other_config param.
            # If we don't do this and we later perform a "convert to template" on this VM,
            # the GUI will infer from the presence of this key that it must query the user
            # for the install media location.  This is unecessary since the template already
            # contains a fully installed disk image, that only needs to be copied.
            xapi.VM.remove_from_other_config(session, vm_ref, "install-distro")

    finally:
        xapi.session.logout(session)

def main():
    try:
        argv = sys.argv[1:]
        for a in ['--default_args=', '--extra_args=', '--args=']:
            while a in argv:
                argv.remove(a)
        opts, mandargs = getopt.getopt(argv, "q",
            ["vm=", "quiet", "args=", "extra_args=", "default_args="])
    except getopt.GetoptError:
        raise UsageError

    vm = None
    img = None
    args = ""
    for opt, val in opts:
        if opt == "--vm":
            vm = val
        if opt in ["--args", "--extra_args", "--default_args"]:
            args += val + " "

    if len(mandargs) > 1:
        raise UsageError

    img = mandargs[0]

    # support running this bootloader multiple times.  We switch bootloader
    # if all required rounds are completed
    other_config = canonicaliseOtherConfig(vm)
    current_round = int(other_config['install-round'])
    
    # how many rounds are required?
    try:
        distro = distros[other_config['install-distro']]
        rounds_required = rounds[distro]
    except:
        raise UnsupportedInstallMethod, "Distribution '%s' is not supported." % other_config['install-distro']

    if current_round == 1:
        handle_first_boot(vm, img, args, other_config)
    elif current_round == 2:
        handle_second_boot(vm, img, args, other_config)

    update_rounds(vm, current_round, rounds_required)

    return 0

if __name__ == "__main__":
    try:
        sys.exit(main())
    except APILevelException, e:
        ex = sys.exc_info()
        err = traceback.format_exception(*ex)

        # print the error out nicely
        print >>sys.stderr, e.apifmt()

        # now log the traceback to syslog
        for exline in err:
            syslog.syslog(syslog.LOG_USER | syslog.LOG_ERR, exline)
        sys.exit(1)
    except UsageError, e:
        print >> sys.stderr, "Invalid usage. Usage: eliloader --vm <vm> <image>"
        sys.exit(2)
