PARTITION_TYPES = {0: "Empty", 238: "GPT protective MBR"}


class MBR:
    def __init__(self):
        self.bootstrap = None
        self.entries = []
        self.signature = None

    def read_from_file(self, filename):
        f = open(filename, 'rb')

        self.bootstrap = f.read(0x01BE)

        for i in range(4):
            self.entries.append(f.read(16))

        self.signature = f.read(2)

        f.close()

        assert self.signature == bytes([0x55, 0xAA])

    def print_entries(self):
        def extract_n_bits(x, start, n):
            mask = ((1 << x) - 1) << start;
            return n & mask;

        for i, entry in enumerate(self.entries):
            print("==== Entry:", i, "====")
            print("Status:", entry[0])

            print("First sector CHS:")

            cylinder = extract_n_bits(entry[2], 0, 2) + (entry[3] << 8)
            sector = extract_n_bits(entry[2], 2, 6)

            print("* head:", entry[1])
            print("* cylinder:", cylinder)
            print("* sector:", sector)

            print("Partition type:", PARTITION_TYPES[entry[4]])

            print("Last sector CHS:")

            cylinder = extract_n_bits(entry[6], 0, 2) + (entry[7] << 8)
            sector = extract_n_bits(entry[6], 2, 6)

            print("* head:", entry[5])
            print("* cylinder:", cylinder)
            print("* sector:", sector)

            print("LBA of first sector:", hex(int.from_bytes(entry[8:12], 'little')))

            print("Number of sectors:", int.from_bytes(entry[12:16], 'little'))





mbr = MBR()

mbr.read_from_file('sda')
#mbr.read_from_file('/dev/sda')