vmdk-raw-parts

#!/usr/bin/env python
import errno
import math
import os
from optparse import OptionParser
import pipes
import random
import re
import string
import struct
import sys
from UserString import MutableString
from uuid import UUID
from collections import namedtuple

# Supports only GPT and MBR partitioning
class Partitioning(object):
    BLOCK = 512
    GPT_SIGNATURE = "EFI PART"
    GPT_VERSION = 0x010000
    MBR_PARTS_OFFSET = 0x1be
    MAX_POST_MBR = 4096

    class FormatError(Exception):
        pass

    class GPTPart(object):
        def __init__(self, num, data):
            self.num, self.data = num, data
            ptype, uuid, start, end, flags, name = \
                struct.unpack_from('<16s16sQQQ72s', data)
            self.type = UUID(bytes_le = ptype)
            if not self.valid(): # unused
                return

            self.uuid = UUID(bytes_le = uuid)
            self.offset = start
            self.size = end - start + 1
            self.flags = flags
            # null terminated
            self.name = name.decode('utf_16_le').split("\0")[0]

        def valid(self):
            return self.type.int != 0
        def __str__(self):
            return ('<GPTPart num=%(num)d name="%(name)s" type=%(type)s ' +
                'uuid=%(uuid)s offset=%(offset)d size=%(size)d ' +
                'flags=%(flags)d>') % self.__dict__

    class MBRPart(object):
        def __init__(self, num, data):
            self.num, self.data = num, data
            status, self.type, offset, size = \
                struct.unpack('<B3xB3xLL', data)
            self.offset, self.size = offset, size
            self.active = (status == 0x80)
            if not self.valid():
                return
        def valid(self):
            return self.type != 0
        def extended(self):
            return self.type in (0x05, 0x0f, 0x85)
        def ebr(self, offset, data, num):
            self.ebr_offset = offset
            self.ebr_data = data
            self.offset += offset
            self.num = num
        def __str__(self):
            return ('<MBRPart num=%(num)d active=%(active)-5s ' +
                'type=%(type)02x offset=%(offset)d size=%(size)d>'
                ) % self.__dict__

    def dump(self):
        parts = list(self.mbr_parts)
        if self.is_gpt():
            parts += self.gpt_parts
        for p in parts:
            if p.valid():
                print p

    def parts(self):
        parts = getattr(self, 'gpt_parts', [])
        if not parts:
            parts = getattr(self, 'mbr_parts')
        return [p for p in parts if p.valid()]

    def extents(self):
        ContentExtent = namedtuple('ContentExtent', 'offset size name contents')
        PartExtent = namedtuple('PartExtent', 'offset size num')

        yield ContentExtent(0, 1, 'mbr', self.mbr)
        parts = self.parts()
        if self.is_gpt():
            yield ContentExtent(1, 1, 'gpt_header', self.gpt_header)
            entries_size = len(self.gpt_entries) / Partitioning.BLOCK
            yield ContentExtent(2, entries_size, 'gpt_entries',
                self.gpt_entries)
        elif self.post_mbr:
            yield ContentExtent(1, len(self.post_mbr) / Partitioning.BLOCK,
                'post_mbr', self.post_mbr)

        for part in parts:
            if getattr(part, 'ebr_offset', None):
                yield ContentExtent(part.ebr_offset, 1, 'ebr' + str(part.num),
                    part.ebr_data)
            yield PartExtent(part.offset, part.size, part.num)

        if self.is_gpt():
            yield ContentExtent(self.gpt_alternate_lba - entries_size,
                entries_size, 'gpt_entries', None)
            yield ContentExtent(self.gpt_alternate_lba, 1, 'gpt_alternate',
                self.gpt_alternate)

    def is_gpt(self):
        return getattr(self, 'gpt_parts', None)

    def mbr_activated(self, active = None):
        buf = MutableString(self.mbr)
        for i in xrange(0, 4):
            off = Partitioning.MBR_PARTS_OFFSET + 16 * i
            buf[off] = struct.pack('B', 0x80 if i + 1 == active else 0)
        return str(buf)

    def mbr_raw(self, data):
        for i in xrange(0, 4):
            off = Partitioning.MBR_PARTS_OFFSET + 16 * i
            yield Partitioning.MBRPart(i + 1, data[off:off+16])

    def mbr_read(self, disk):
        self.mbr_parts = []

        ext = None
        for p in self.mbr_raw(self.mbr):
            if p.extended():
                ext = p.offset
            elif p.valid():
                self.mbr_parts.append(p)

        if ext:
            pos = ext
            n = 5
            while True:
                disk.seek(pos * Partitioning.BLOCK)
                data = disk.read(Partitioning.BLOCK)
                parts = list(self.mbr_raw(data))
                if parts[2].type != 0 or parts[3].type != 0:
                    raise FormatError()

                parts[0].ebr(pos, data, n)
                self.mbr_parts.append(parts[0])

                if parts[1].type == 0:
                    break
                pos = ext + parts[1].offset
                n += 1
            disk.seek(1)

    def gpt_read(self, disk):
        self.gpt_header = disk.read(Partitioning.BLOCK)
        sig, vers, self.gpt_alternate_lba, first_usable, pcount, psize = \
            struct.unpack_from('<8sL20xQQ32xLL', self.gpt_header)
        if sig != Partitioning.GPT_SIGNATURE or \
                vers != Partitioning.GPT_VERSION:
            self.gpt_parts = None
            return False

        self.gpt_entries = disk.read((first_usable - 2) * Partitioning.BLOCK)
        self.gpt_parts = []
        for i in xrange(0, pcount):
            part = Partitioning.GPTPart(i + 1,
                self.gpt_entries[i * psize : (i + 1) * psize])
            if part.valid():
                self.gpt_parts.append(part)

        disk.seek(self.gpt_alternate_lba * Partitioning.BLOCK)
        self.gpt_alternate = disk.read(Partitioning.BLOCK)
        return True

    def __init__(self, diskpath, options):
        disk = open(diskpath)
        try:
            self.mbr = disk.read(Partitioning.BLOCK)
            if self.mbr[510:512] != "\x55\xAA":
                raise FormatError()

            self.mbr_read(disk)
            if options.mbr or not self.gpt_read(disk):
                # Include small post-MBR bootloader space
                self.post_mbr = None
                if self.mbr_parts[0].offset < Partitioning.MAX_POST_MBR:
                    disk.seek(Partitioning.BLOCK)
                    self.post_mbr = disk.read((self.mbr_parts[0].offset - 1) *
                        Partitioning.BLOCK)
        finally:
            disk.close()

class PartitionedVMDK(object):
    class PartSpecError(Exception):
        pass

    def __init__(self, diskpath, options):
        self.diskpath = os.path.realpath(diskpath)
        self.platform = options.platform
        if self.platform == None:
            self.platform = sys.platform
        self.diskparts = Partitioning(self.diskpath, options)
        self.drop_privileges()

    def drop_privileges(self):
        if os.getuid() == 0:
            olduid = os.environ.get('SUDO_UID')
            if olduid:
                olduid = int(olduid)
                oldgid = int(os.environ['SUDO_GID'])
                os.setgroups([])
                os.setgid(oldgid)
                os.setuid(olduid)

    def writefile(self, dir, name, contents):
        f = open(os.path.join(dir, name), 'w')
        f.write(contents)
        f.close()

    def write_script(self, dir, args):
        if args == None:
            return

        script = "#!/bin/sh\ncd $parent\n$sudo $command $args\n"
        script = string.Template(script).substitute(
            parent = pipes.quote(os.path.realpath(os.path.dirname(dir))),
            command = pipes.quote(os.path.realpath(__file__)),
            args = str.join(' ', [pipes.quote(a) for a in args]),
            sudo = ('' if os.access(self.diskpath, os.R_OK) else 'sudo'))
        name = 'regen.sh'
        self.writefile(dir, name, script)
        os.chmod(os.path.join(dir, name), 0770)

    def write(self, dir, wantparts, active = None, args = None):
        wanted = self.wanted(wantparts)
        if active:
            active_gpt = self.find_part(active)
            active = [p for p in self.diskparts.mbr_parts
                if p.offset == active_gpt.offset][0].num

        try:
            os.mkdir(dir)
        except OSError as e:
            if e.errno != errno.EEXIST:
                raise

        for ext in self.diskparts.extents():
            if getattr(ext, 'contents', None):
                contents = ext.contents
                if ext.name == 'mbr':
                    contents = self.diskparts.mbr_activated(active)
                self.writefile(dir, ext.name, contents)

        base = os.path.basename(dir)
        base, _ = os.path.splitext(base)
        self.writefile(dir, '%s.vmdk' % base, self.vmdk(wanted))
        self.write_script(dir, args)

    def find_part(self, key):
        parts = self.diskparts.parts()
        prop = 'name'
        try:
            key = int(key)
            prop = 'num'
        except ValueError: # Not an integer
            pass

        return next(p for p in parts if getattr(p, prop, None) == key)

    # return an ordered list of numbers of desired partitions
    def wanted(self, want):
        return sorted([self.find_part(w).num for w in want])

    def macpart(self, partnum):
        return re.sub(r'/r?disk(\d+)$', r'/disk\1s' + str(partnum),
            self.diskpath)

    def extent(self, elist, size, partnum = None, offset = None, name = None):
        if name:
            spec = 'FLAT "%s" 0' % name
        elif partnum and self.platform == 'darwin':
            spec = 'FLAT "%s" 0' % self.macpart(partnum)
        elif offset and self.platform != 'darwin':
            spec = 'FLAT "%s" %d' % (self.diskpath, offset)
        else:
            spec = 'ZERO'
        ext = 'RW %d %s' % (size, spec)
        elist.append(ext)

    def extents(self, wanted):
        desc = []
        exts = sorted(self.diskparts.extents(), key = lambda e: e.offset)

        pos = 0
        for ext in exts:
            if ext.offset > pos: # Add some zeros
                self.extent(desc, ext.offset - pos)

            if hasattr(ext, 'name'):
                self.extent(desc, ext.size, name = ext.name)
            elif ext.num in wanted:
                self.extent(desc, ext.size, offset = ext.offset,
                    partnum = ext.num)
            else:
                self.extent(desc, ext.size)
            pos = ext.offset + ext.size

        return str.join("\n", desc)

    def vmdk(self, wanted):
    	# Need something in DDB or VBox complains
        template = """# Disk DescriptorFile
version=1
encoding="UTF-8"
CID=$cid
parentCID=ffffffff
isNativeSnapshot="no"
createType="partitionedDevice"

# Extent description
$extents

ddb.virtualHWVersion = "4"
"""
        cid = "%08x" % random.randint(0, 2**32 - 1)
        extents = self.extents(wanted)
        return string.Template(template).substitute(cid = cid,
            extents = extents)

if __name__ == "__main__":
    orig_args = sys.argv[1:]
    parser = OptionParser(
    	usage = 'Usage: %prog [options] DEVICE OUTPUT [PART1 ...]',
    	description =
"""Creates a virtual disk that exposes to one or more partitions of
a physical device to the VM.

The DEVICE must have an MBR or GPT partition table. On successful exit,
OUTPUT will be a directory containing a VMDK disk and support files,
which should work in either VMWare or VirtualBox. The partitions to
expose are specified by integers, starting at one; or by name.""")
    parser.add_option('-b', '--boot', metavar = 'PART', dest = "boot",
    	help = "Make this partition bootable")
    parser.add_option('-m', '--mbr', dest = 'mbr', action = 'store_true',
    	help = "Ignore any GPT partitioning")
    parser.add_option('-p', '--platform', dest = "platform",
    	help = "Target the given platform (darwin or linux)")
    parser.add_option('-v', '--verbose', dest = "verbose",
    	help = "Dump the partition table as we run", action = 'store_true')
    (options, args) = parser.parse_args(orig_args)
    if len(args) < 2:
        parser.print_help()
        sys.exit(-1)
    disk = args.pop(0)
    dest = os.path.normpath(args.pop(0))
    parts = args

    vmdk = PartitionedVMDK(disk, options)
    if options.verbose:
        vmdk.diskparts.dump()
    vmdk.write(dest, parts, active = options.boot, args = orig_args)