mbr_parser.py

#!/usr/bin/env python3
import sys
import struct

# https://www.win.tue.nl/~aeb/partitions/partition_types-1.html
SYSTEM_IDS = {
    0x00: "Empty",
    0x07: "NTFS / exFAT",
    0x82: "Linux swap",
    0x83: "Linux native partiton",
    0xEE: "EFI header",
    0xEF: "EFI Filesystem"
}

class MBRPartitionEntry:
    ENTRY_STRUCT = struct.Struct("BBHBBHII")

    def __init__(self, attributes, chs_start, partition_type, chs_end, lba_start, num_sectors):
        self.attributes = attributes
        self.chs_start = chs_start
        self.partition_type = partition_type
        self.chs_end = chs_end
        self.lba_start = lba_start
        self.num_sectors = num_sectors

    @property
    def bootable(self):
        # bit 7 set = Drive is bootable
        return bool(self.attributes & (1 << 7))

    @property
    def partition_type_name(self):
        return SYSTEM_IDS.get(self.partition_type, "Unknown")

    @staticmethod
    def _parse_chs(head, sector_and_cylinder):
        # Upper 10 bits are the cylinder
        cylinder = (sector_and_cylinder & 0xFFC0) >> 6
        # Lower 6 bits are the sector
        sector = sector_and_cylinder & 0x3F

        return (cylinder, head, sector)

    @classmethod
    def from_raw(cls, data):
        unpacked = cls.ENTRY_STRUCT.unpack(data)

        attributes = unpacked[0]
        chs_start = cls._parse_chs(unpacked[1], unpacked[2])
        partition_type = unpacked[3]
        chs_end = cls._parse_chs(unpacked[4], unpacked[5])

        return cls(attributes, chs_start, partition_type, chs_end, unpacked[6], unpacked[7])


class MBRPartitionTable(tuple):
    ENTRY_SIZE = 16
    MAX_ENTRIES = 4
    TABLE_START = 0x1BE
    TABLE_END = TABLE_START + ENTRY_SIZE * MAX_ENTRIES

    @classmethod
    def from_raw(cls, data):
        partition_range = range(cls.TABLE_START, cls.TABLE_END, cls.ENTRY_SIZE)

        entries = []
        for raw_entry in [data[i:i+cls.ENTRY_SIZE] for i in partition_range]:
            # Ignore empty partitions
            if raw_entry == b'\x00' * cls.ENTRY_SIZE:
                entries.append(None)
                continue

            entry = MBRPartitionEntry.from_raw(raw_entry)
            entries.append(entry)

        return cls(entries)


def main():
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <mbr-file>")
        return

    mbr_filename = sys.argv[1]
    with open(mbr_filename, "rb") as mbr_file:
        raw_mbr = mbr_file.read(512)
    
    if raw_mbr[-2:] != b'\x55\xAA':
        print("File is not a valid Master Boot Record (does not contain signature)")
        return

    mbr = MBRPartitionTable.from_raw(raw_mbr)

    print("-------------[ MBR Partition Table ]-------------\n")
    for i, entry in enumerate(mbr):
        if entry is None:
            print(f"[ Partition {i} ]\nUnused entry\n")
        else:
            print(f"[ Partition {i} ]\n"
                  f"Bootable\t: {entry.bootable}\n"
                  f"Partition tyoe\t: {entry.partition_type_name} ({entry.partition_type:#02x})\n"
                  f"Start CHS\t: {entry.chs_start}\n"
                  f"End CHS  \t: {entry.chs_end}\n"
                  f"Start LBA\t: {entry.lba_start}\n"
                  f"Num of Sectors\t: {entry.num_sectors}\n")


if __name__ == "__main__":
    main()