fix: write partition data only

Fixes an issue where the partition data would overwrite the entire disk.
swysocki · Sep 7, 2023 · e82041e · e82041e
1 parent 1bb9ef2
commit e82041e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/gpt_image/partition.py b/gpt_image/partition.py
@@ -195,7 +195,7 @@ def write_data(self, disk: Disk, data: bytes) -> int:
         """
         if len(data) > self.size:
             raise ValueError(f"data too large for partition: {len(data)} {self.size}")
-        with open(str(disk.image_path), "wb") as b:
+        with open(str(disk.image_path), "r+b") as b:
             start = self.first_lba * disk.sector_size
             b.seek(start)
             b.write(data)

diff --git a/tests/test_disk.py b/tests/test_disk.py
@@ -67,12 +67,12 @@ def test_write_data(new_image):
         byte_count = len(BYTE_DATA)
         b.seek(start)
         test_buffer = b.read(byte_count)
-        assert test_buffer == BYTE_DATA
+        assert bytes(test_buffer) == BYTE_DATA
 
 
 def test_read_partition(new_image):
     disk = Disk.open(new_image)
     part = disk.table.partitions.find("partition1")
     part.write_data(disk, BYTE_DATA)
     read_data = part.read(disk)
-    assert read_data == BYTE_DATA
+    assert read_data[:len(BYTE_DATA)] == BYTE_DATA