diff --git a/gpt/__init__.py b/gpt/__init__.py index d52a765..789c575 100644 --- a/gpt/__init__.py +++ b/gpt/__init__.py @@ -142,7 +142,7 @@ def calculate_header_crc32(self): self.number_of_partition_entries, self.size_of_partition_entry, self.partition_entry_array_crc32) - return binascii.crc32(header_crc32_input) + return binascii.crc32(header_crc32_input) & 0xffffffff class GPTPartitionEntry(): @@ -172,22 +172,22 @@ def is_empty(self): def calculate_partition_entry_array_crc32(data): - return binascii.crc32(data) + return binascii.crc32(data) & 0xffffffff def encode_mbr(mbr): - partition_record = bytes() + partition_record = bytearray() for i in range(0, 4): - partition = input.partitions[i] - partition_record.append(pack('< B B B B B B B B I I', + partition = mbr.partitions[i] + partition_record.extend(pack('< B B B B B B B B I I', partition.boot_indicator, - partition.start_chs[0], partition.start_chs[1], partition.start_chs[2], + partition.start_chs[0], partition.os_type, - partition.end_chs[0], partition.end_chs[1], partition.end_chs[2], + partition.end_chs[0], partition.lba_ss[0], partition.lba_ss[1])) @@ -195,7 +195,7 @@ def encode_mbr(mbr): mbr.bootstrap_code, mbr.unique_mbr_disk_signature, mbr.unknown, - partition_record, + bytes(partition_record), mbr.signature) return output @@ -298,9 +298,9 @@ def encode_gpt_partition_entry(gpt_partition_entry): data = pack('<16s 16s Q Q Q 72s', gpt_partition_entry.partition_type_guid_raw, gpt_partition_entry.unique_partition_guid_raw, - gpt_partition_entry.start_lba, + gpt_partition_entry.starting_lba, gpt_partition_entry.ending_lba, - gpt_partition_entry.attributes, + gpt_partition_entry.attributes_raw, gpt_partition_entry.partition_name_raw) return data @@ -322,14 +322,14 @@ def decode_gpt_partition_entry(data): def encode_gpt_partition_entry_array(gpt_partition_entries, size, count): - data = bytes() + data = bytearray() for i in range(0, count): d = encode_gpt_partition_entry(gpt_partition_entries[i]) - data.append(d) + data.extend(d) # fill with zeroes if less than size if len(d) < size: - data.append(bytes(size - len(d))) - return data + data.extend(bytearray(size - len(d))) + return bytes(data) def decode_gpt_partition_entry_array(data, size, count): diff --git a/setup.cfg b/setup.cfg index fc0cddd..3feb1f9 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,9 +26,14 @@ classifiers = zip_safe = False packages = gpt include_package_data = True +setup_requires = pytest-runner +tests_require = pytest [options.entry_points] console_scripts = print_mbr=gpt.scripts:print_mbr print_gpt_header=gpt.scripts:print_gpt_header print_gpt_partition_entry_array=gpt.scripts:print_gpt_partition_entry_array + +[aliases] +test=pytest diff --git a/tests/sample-data/disk1.bin b/tests/sample-data/disk1.bin new file mode 100644 index 0000000..07c43a7 Binary files /dev/null and b/tests/sample-data/disk1.bin differ diff --git a/tests/sample-data/disk2.bin b/tests/sample-data/disk2.bin new file mode 100644 index 0000000..c6dc714 Binary files /dev/null and b/tests/sample-data/disk2.bin differ diff --git a/tests/sample-data/merged.bin b/tests/sample-data/merged.bin new file mode 100644 index 0000000..5d9e681 Binary files /dev/null and b/tests/sample-data/merged.bin differ diff --git a/tests/test_gpt.py b/tests/test_gpt.py new file mode 100644 index 0000000..0da8f8e --- /dev/null +++ b/tests/test_gpt.py @@ -0,0 +1,113 @@ +import pytest +import os + + +def load_sample_data(name): + path = os.path.join(os.path.dirname(__file__), + 'sample-data', '%s.bin' % (name,)) + with open(path, 'rb') as f: + return f.read() + + +def entries_range(h): + start = h.partition_entry_lba * 0x200 + length = h.number_of_partition_entries * h.size_of_partition_entry + end = start + length + return start, end + + +def test_roundtrip_mbr(): + import gpt + sample = load_sample_data('disk1') + raw1 = sample[:0x200] + m = gpt.decode_mbr(raw1) + assert m.is_valid() + raw2 = gpt.encode_mbr(m) + assert raw1 == raw2 + + +def test_roundtrip_gpt_header(): + import gpt + sample = load_sample_data('disk1') + raw1 = sample[0x200:0x400] + h = gpt.decode_gpt_header(raw1) + assert h.header_size == 92, "expected header size 92" + raw2 = gpt.encode_gpt_header(h) + assert raw1[:h.header_size] == raw2 + + +def test_roundtrip_gpt_partition_entry_array(): + import gpt + sample = load_sample_data('disk1') + h = gpt.decode_gpt_header(sample[0x200:0x400]) + + start, end = entries_range(h) + + decoded_entries = gpt.decode_gpt_partition_entry_array( + sample[start:end], + h.size_of_partition_entry, + h.number_of_partition_entries) + + reencoded_entries = gpt.encode_gpt_partition_entry_array( + decoded_entries, + h.size_of_partition_entry, + h.number_of_partition_entries) + + assert sample[start:end] == reencoded_entries + + +def test_merge(): + import gpt + + sample1 = load_sample_data('disk1') + sample2 = load_sample_data('disk2') + sample_merged = load_sample_data('merged') + + h1 = gpt.decode_gpt_header(sample1[0x200:0x400]) + h2 = gpt.decode_gpt_header(sample2[0x200:0x400]) + + start1, end1 = entries_range(h1) + start2, end2 = entries_range(h2) + + decoded_entries1 = gpt.decode_gpt_partition_entry_array( + sample1[start1:end1], + h1.size_of_partition_entry, + h1.number_of_partition_entries) + + decoded_entries2 = gpt.decode_gpt_partition_entry_array( + sample2[start2:end2], + h2.size_of_partition_entry, + h2.number_of_partition_entries) + + # Copy disk2 partition table + merged_h = gpt.decode_gpt_header(sample2[0x200:0x400]) + merged_entries = list(decoded_entries2) + + # Copy partitions 1-4 from disk1 + merged_entries[0:4] = decoded_entries1[0:4] + merged_entries_raw = gpt.encode_gpt_partition_entry_array( + merged_entries, + merged_h.size_of_partition_entry, + merged_h.number_of_partition_entries) + + # Update crc32 values + merged_h.partition_entry_array_crc32 = gpt.calculate_partition_entry_array_crc32(merged_entries_raw) + merged_h.header_crc32 = merged_h.calculate_header_crc32() + + # Should be valid + assert merged_h.is_valid() + + # Encode header + merged_h_raw = gpt.encode_gpt_header(merged_h) + + m_start, m_end = entries_range(merged_h) + assert m_end - m_start == len(merged_entries_raw) + + # Construct the new partition table + merged_raw = bytearray(0x400) + merged_raw[:0x200] = sample1[:0x200] # Copy protective MBR from disk1 + merged_raw[0x200:0x200+merged_h.header_size] = merged_h_raw + merged_raw[m_start:m_end] = merged_entries_raw + + # Did we get what we expect? + assert merged_raw == sample_merged diff --git a/tox.ini b/tox.ini new file mode 100644 index 0000000..52065f2 --- /dev/null +++ b/tox.ini @@ -0,0 +1,7 @@ +[tox] +envlist = py27,py37 + +[testenv] +commands = pytest +deps = + pytest