#!/usr/bin/python
#
# lspart.py
#
# By Ed Plese <ed@edplese.com>
#
# Parses MBR of the given device and displays the partitions on the device,
# including information such as the offset of the partition in bytes, the
# size of the partition, and its type.
#
# An example of the output is:
# # lspart.py /dev/zvol/dsk/rpool/xvm/win2k8@installed
#   Start Offset    Size  Type
#        1048576  100.0M  07 Windows NTFS
#      105906176   15.9G  07 Windows NTFS
#              0    0.0B  00 Empty
#              0    0.0B  00 Empty

import struct
import sys

SECTOR_SIZE = 512

# Adapted from http://www.win.tue.nl/~aeb/partitions/partition_types-1.html
TYPE_LOOKUP = {
    0x00: "Empty",
    0x04: "DOS FAT16",
    0x05: "Extended Partition",
    0x06: "DOS FAT16",
    0x07: "Windows NTFS",
    0x08: "AIX",
    0x0b: "Windows FAT32",
    0x0c: "Windows FAT32",
    0x0e: "Windows FAT16",
    0x0f: "Windows Extended Partition",
    0x27: "Windows RE Hidden Partition",
    0x35: "JFS",
    0x42: "Windows 2000 Dynamic Disk",
    0x82: "Linux Swap / Solaris",
    0x83: "Linux Native Partition",
    0x85: "Linux Extended Partition",
    0x86: "Linux RAID / FAT16 Volume Set",
    0x87: "NTFS Volume Set",
    0x8e: "Linux LVM",
    0x9f: "BSD",
    0xa5: "BSD",
    0xa6: "OpenBSD",
    0xa8: "Mac OS-X",
    0xa9: "Net-BSD",
    0xab: "Mac OS-X Boot",
    0xb7: "BSD",
    0xb8: "BSD Swap",
    0xbe: "Solaris 8 Boot",
    0xbf: "Solaris",
    0xfb: "VMware",
    0xfc: "VMware",
    0xfd: "Linux RAID",
    }

class Volume:
    def __init__(self, path):
        self.f = open(path, "rb")

    def read(self, address, size):
        self.f.seek(address)
        return self.f.read(size)

    def printBootRecords(self, address=0, level=0):
        br = BootRecord(self.read(address, 512))
        for partition in br:
            print "%14d  %04s  %s%02x %s" % (address + partition.lba * SECTOR_SIZE,
                human_size(partition.blocks * SECTOR_SIZE),
                "   " * level, partition.type,
                TYPE_LOOKUP.get(partition.type, "Unknown"))

            if partition.type == 0x05:
                self.printBootRecords(address + partition.lba * SECTOR_SIZE,
                level + 1)

class BootRecord:
    def __init__(self, data):
        assert self.hasValidSignature(data)
        self.table = data[0x01be:0x01be + 64]

    def hasValidSignature(self, data):
        return data[510:512] == "\x55\xaa"

    def __iter__(self):
        for i in range(4):
            yield Partition(self.table[i * 16:i * 16 + 16])

class Partition:
    def __init__(self, data):
        fmt = "8B2I"
        fields = struct.unpack(fmt, data)

        self.type = fields[4]
        self.lba = fields[8]
        self.blocks = fields[9]

def human_size(bytes):
    b = float(bytes)
    for prefix in "BKMGTPEZY":
        if b < 1000.0:
            return "%5.1f%s" % (b, prefix)
        b /= 1024.0

    return str(bytes)

def usage():
    print "Usage: lspart.py <dev>"
    sys.exit(1)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        usage()

    try:
        v = Volume(sys.argv[1])
    except:
        print "Unable to open file or device:", sys.argv[1]
        usage()
    print "  Start Offset    Size  Type"
    v.printBootRecords()