xen/domUloader.py

#!/usr/bin/env python
# domUloader.py
"""Loader for kernel and (optional) ramdisk from domU filesystem

Given a physical disk (or disk image) for a domU and the path of a kernel and
optional ramdisk, copies the kernel and ramdisk from the domU disk to a
temporary location in dom0.

The --entry parameter specifies the location of the kernel (and optional
ramdisk) within the domU filesystem.  dev is the disk as seen by domU.
Filenames are relative to that filesystem.

The disk is passed as the last parameter.  It must be a block device or raw
disk image.  More complex disk images (QCOW, VMDK, etc) must already be
configured via blktap and presented as a block device.

The script writes an sxpr specifying the locations of the copied kernel and
ramdisk into the file specified by --output (default is stdout).

Limitations:
 - It is assumed both kernel and ramdisk are on the same filesystem.
 - domUs might use LVM; the script currently does not have support for setting
   up LVM mappings for domUs; it's not trivial and we might risk namespace
   conflicts. If you want to use LVM inside domUs, set up a small non-LVM boot
   partition and specify it in bootentry.
 
The script uses kpartx (multipath-tools) to create mappings for devices that
are exported as whole disk devices that are partitioned.

(c) 01/2006 Novell Inc
License: GNU GPL
Author: Kurt Garloff <garloff@suse.de>
"""

import os, sys, getopt
from stat import *
from xen.xend import sxp
import tempfile
import time

# Global options
quiet = False
verbose = False
dryrun = False
tmpdir = '/var/lib/xen/tmp'

# Helper functions

def error(s):
    print >> sys.stderr, "domUloader error: %s" % s

def verbose_print(s):
    if verbose:
        print >> sys.stderr, "domUloader: %s" % s

def traildigits(strg):
    """Return the trailing digits, used to split the partition number off"""
    idx = len(strg)-1
    while strg[idx].isdigit():
        if len == 0:
            return strg
        idx -= 1
    return strg[idx+1:]    

def getWholedisk(part):
    while len(part) and part[len(part)-1].isdigit():
        part = part[:-1]
    return part

#def isWholedisk(domUname):
#    """Determines whether dev is a wholedisk dev"""
#    return not domUname[-1:].isdigit()


class Wholedisk:
    "Class representing a whole disk that may have partitions"
    def __init__(self, vdev, pdev):
        "c'tor: set up"
        self.is_blk = (S_ISBLK(os.stat(pdev)[ST_MODE]))
        self.ldev = None
        self.vdev = vdev
        self.pdev = pdev
        self.mapped = 0
        self.partitions = []
        self.pcount = self.scanpartitions()

    def physdev(self):
        """Gets the physical device used to access the device from dom0"""
        if self.ldev:
            return self.ldev
        return self.pdev

    def findPart(self, vdev):
        "Find device dev in list of partitions"
        if len(vdev) > 5 and vdev[:5] == "/dev/":
            vdev = vdev[5:]
        for part in self.partitions:
            if vdev == part.vdev:
                return part
        return None

    def loopsetup(self):
        """Sets up the loop mapping for a disk image.

        Will raise if no loopbacks are available.
        """
        if not self.is_blk and not self.ldev:
            # Loops through all loopback devices, attempting to
            # find a free one to set up.  Don't scan for free and
            # then try to set it up as a separate step - too racy!
            i = 0
            while True:
                ldev = '/dev/loop%i' % (i)
                if not os.path.exists(ldev):
                    break
                i += 1
                fd = os.popen("losetup %s %s 2> /dev/null" % (ldev, self.pdev))
                if not fd.close():
                    verbose_print("losetup %s %s" % (ldev, self.pdev))
                    self.ldev = ldev
                    break
            if not self.ldev:
                raise RuntimeError("No free loop device found")

    def loopclean(self):
        """Delete the loop mapping.

        Will never raise.
        """
        if self.ldev:
            verbose_print("losetup -d %s" % self.ldev)
            # Even seemingly innocent queries like "losetup /dev/loop0"
            # can temporarily block the loopback and cause transient
            # failures deleting the loopback, hence the retry logic.
            retries = 10
            while retries:
                fd = os.popen("losetup -d %s" % self.ldev)
                if not fd.close():
                    self.ldev = None
                    break
                else:
                    time.sleep(0.1)
                    retries -= 1
    
    def scanpartitions(self):
        """Scan device for partitions (kpartx -l) and set up data structures,
           Returns number of partitions found."""
        self.loopsetup()
        # TODO: We could use fdisk -l instead and look at the type of
        #  partitions; this way we could also detect LVM and support it.
        fd = os.popen("kpartx -l %s" % self.physdev())
        pcount = 0
        for line in fd.readlines():
            line = line.strip()
            verbose_print("kpartx -l: %s" % (line,))
            (pname, params) = line.split(':')
            pno = int(traildigits(pname.strip()))
            #if pname.rfind('/') != -1:
            #    pname = pname[pname.rfind('/')+1:]
            #pname = self.pdev[:self.pdev.rfind('/')] + '/' + pname
            pname = "/dev/mapper/" + pname
            verbose_print("Found partition: vdev %s, pdev %s" % ('%s%i' % (self.vdev, pno), pname))
            self.partitions.append(Partition(self, '%s%i' % (self.vdev, pno), pname))
            pcount += 1
        fd.close()
        if not pcount:
            if self.ldev:
                ref = self
            else:
                ref = None
            self.partitions.append(Partition(ref, self.vdev, self.pdev))
        return pcount

    def activatepartitions(self):
        "Set up loop mapping and device-mapper mappings"
        if not self.mapped:
            self.loopsetup()
            if self.pcount:
                verbose_print("kpartx -a %s" % self.physdev())
                fd = os.popen("kpartx -a %s" % self.physdev())
                fd.close()
        self.mapped += 1
    
    def deactivatepartitions(self):
        """Remove device-mapper mappings and loop mapping.

        Will never raise.
        """
        if not self.mapped:
            return
        self.mapped -= 1
        if not self.mapped:
            if self.pcount:
                verbose_print("kpartx -d %s" % self.physdev())
                fd = os.popen("kpartx -d %s" % self.physdev())
                fd.close()
            self.loopclean()

    def __del__(self):
        "d'tor: clean up"
        self.deactivatepartitions()
        self.loopclean()

    def __repr__(self):
        "string representation for debugging"
        strg = "[" + self.vdev + "," + self.pdev + ","
        if self.ldev:
            strg += self.ldev
        strg += "," + str(self.pcount) + ",mapped %ix]" % self.mapped
        return strg

class Partition:
    """Class representing a domU filesystem (partition) that can be
       mounted in dom0"""
    def __init__(self, whole = None, vdev = None, pdev = None):
        "c'tor: setup"
        self.wholedisk = whole
        self.vdev    = vdev
        self.pdev    = pdev
        self.mountpoint = None

    def __del__(self):
        "d'tor: cleanup"
        if self.mountpoint:
            self.umount()
        # Not needed: Refcounting will take care of it.
        #if self.wholedisk:
        #    self.wholedisk.deactivatepartitions()

    def __repr__(self):
        "string representation for debugging"
        strg = "[" + self.vdev + "," + self.pdev + ","
        if self.mountpoint:
            strg += "mounted on " + self.mountpoint + ","
        else:
            strg += "not mounted,"
        if self.wholedisk:
            return strg + self.wholedisk.__repr__() + "]"
        else:
            return strg + "]"

    def mount(self, fstype = None, options = "ro"):
        "mount filesystem, sets self.mountpoint"
        if self.mountpoint:
            return
        if self.wholedisk:
            self.wholedisk.activatepartitions()
        mtpt = tempfile.mkdtemp(prefix = "%s." % self.vdev, dir = tmpdir)
        mopts = ""
        if fstype:
            mopts += " -t %s" % fstype
        mopts += " -o %s" % options
        verbose_print("mount %s %s %s" % (mopts, self.pdev, mtpt))
        fd = os.popen("mount %s %s %s" % (mopts, self.pdev, mtpt))
        err = fd.close()
        if err:
            raise RuntimeError("Error %i from mount %s %s on %s" % \
                (err, mopts, self.pdev, mtpt))
        self.mountpoint = mtpt
        
    def umount(self):
        """umount filesystem at self.mountpoint"""
        if not self.mountpoint:
            return
        verbose_print("umount %s" % self.mountpoint)
        fd = os.popen("umount %s" % self.mountpoint)
        err = fd.close()
        try:
            os.rmdir(self.mountpoint)
        except:
            pass
        if err:
            error("Error %i from umount %s" % (err, self.mountpoint))
        else:
            self.mountpoint = None
        if self.wholedisk:
            self.wholedisk.deactivatepartitions()

def parseEntry(entry):
    "disects bootentry and returns vdev, kernel, ramdisk"
    def bad():
        raise RuntimeError, "Malformed --entry"
    fsspl = entry.split(':')
    if len(fsspl) != 2:
        bad()
    vdev = fsspl[0]
    entry = fsspl[1]
    enspl = entry.split(',')
    if len(enspl) not in (1, 2):
        bad()
    # Prepend '/' if missing
    kernel = enspl[0]
    if kernel == '':
        bad()
    if kernel[0] != '/':
        kernel = '/' + kernel
    ramdisk = None
    if len(enspl) > 1:
        ramdisk = enspl[1]
        if ramdisk != '' and ramdisk[0] != '/':
            ramdisk = '/' + ramdisk
    return vdev, kernel, ramdisk

def copyFile(src, dst):
    "Wrapper for shutil.filecopy"
    import shutil
    verbose_print("cp %s %s" % (src, dst))
    stat = os.stat(src)
    if stat.st_size > 16*1024*1024:
        raise RuntimeError("Too large file %s (%s larger than 16MB)" \
            % (src, stat.st_size)) 
    try:
        shutil.copyfile(src, dst)
    except:
        os.unlink(dst)
        raise()

def copyKernelAndRamdisk(disk, vdev, kernel, ramdisk):
    """Finds vdev in list of partitions, mounts the partition, copies
       kernel [and ramdisk] off to dom0 files, umounts the parition again,
       and returns sxpr pointing to these copies."""
    verbose_print("copyKernelAndRamdisk(%s, %s, %s, %s)" % (disk, vdev, kernel, ramdisk))
    if dryrun:
        return "linux (kernel kernel.dummy) (ramdisk ramdisk.dummy)"
    part = disk.findPart(vdev)
    if not part:
        raise RuntimeError("Partition '%s' does not exist" % vdev)
    part.mount()
    try:
        (fd, knm) = tempfile.mkstemp(prefix = "kernel.", dir = tmpdir)
        os.close(fd)
        copyFile(part.mountpoint + kernel, knm)
    except: 
        os.unlink(knm)
        part.umount()
        raise
    if not quiet:
        print "Copy kernel %s from %s to %s for booting" % (kernel, vdev, knm)
    sxpr = "linux (kernel %s)" % knm
    if ramdisk:
        try:
            (fd, inm) = tempfile.mkstemp(prefix = "ramdisk.", dir = tmpdir)
            os.close(fd)
            copyFile(part.mountpoint + ramdisk, inm)
        except:
            os.unlink(knm)
            os.unlink(inm)
            part.umount()
            raise
        sxpr += "(ramdisk %s)" % inm
    part.umount()
    return sxpr

def main(argv):
    "Main routine: Parses options etc."
    global quiet, dryrun, verbose, tmpdir
    def usage():
        "Help output (usage info)"
        global verbose, quiet, dryrun
        print >> sys.stderr, "domUloader usage: domUloader [--output=fd] [--quiet] [--dryrun] [--verbose]\n" +\
                             "[--help] --entry=dev:kernel[,ramdisk] physdisk [virtdisk]\n"
        print >> sys.stderr, __doc__

    try:
        (optlist, args) = getopt.gnu_getopt(argv, 'qvh', \
            ('entry=', 'output=', 'tmpdir=', 'help', 'quiet', 'dryrun', 'verbose'))
    except:
        usage()
        sys.exit(1)

    entry = None
    output = None
    pdisk = None
    vdisk = None

    for (opt, oarg) in optlist:
        if opt in ('-h', '--help'):
            usage()
            sys.exit(1)
        elif opt in ('-q', '--quiet'):
            quiet = True
        elif opt in ('-n', '--dryrun'):
            dryrun = True
        elif opt in ('-v', '--verbose'):
            verbose = True
        elif opt == '--output':
            output = oarg
        elif opt == '--entry':
            entry = oarg
        elif opt == '--tmpdir':
            tmpdir = oarg

    verbose_print(str(argv))

    if args:
        if len(args) == 2:
            pdisk = args[1]
        elif len(args) == 3:
            pdisk = args[1]
            vdisk = args[2]
    
    if not entry or not pdisk:
        usage()
        sys.exit(1)
    
    if output is None or output == "-":
        fd = sys.stdout.fileno()
    else:
        fd = os.open(output, os.O_WRONLY)
    
    if not os.access(tmpdir, os.X_OK):
        os.mkdir(tmpdir)
        os.chmod(tmpdir, 0750)

    vdev, kernel, ramdisk = parseEntry(entry)
    if not vdisk:
        vdisk = getWholedisk(vdev)
        verbose_print("vdisk not specified; guessing '%s' based on '%s'" % (vdisk, vdev))
    if not vdev.startswith(vdisk):
        error("Virtual disk '%s' does not match entry '%s'" % (vdisk, entry))
        sys.exit(1)
    disk = Wholedisk(vdisk, pdisk)
    
    r = 0
    try:
        sxpr = copyKernelAndRamdisk(disk, vdev, kernel, ramdisk)
        os.write(fd, sxpr)
    except Exception, e:
        error(str(e))
        r = 1

    for part in disk.partitions:
        part.wholedisk = None
    del disk

    return r

# Call main if called (and not imported)
if __name__ == "__main__":
    r = 1
    try:
        r = main(sys.argv)
    except Exception, e:
        error(str(e))
    sys.exit(r)