Bare GPT working

Protective MBR and GPT functioning.  Partition CRC's are not calculated
correctly when a partition is created.

gpt_image/disk.py

@@ -5,7 +5,7 @@ class Geometry:
     """Geometry of disk image
 
     Attributes:
-        sector_size: statically set to 512 bytes
+        sector_size: typically set to 512 bytes
         total_bytes: disk size in bytes
         total_sectors: number of sectors on the disk
         total_lba: number of logical blocks on the disk
@@ -21,9 +21,9 @@ class Geometry:
         backup_header_array_byte: byte where the backup partition array starts
     """
 
-    def __init__(self, size: int) -> None:
+    def __init__(self, size: int, sector_size: int = 512) -> None:
         """Init Geometry with size in bytes"""
-        self.sector_size = 512
+        self.sector_size = sector_size
         self.total_bytes = size
         self.total_sectors = int(size / self.sector_size)
         self.total_lba = int(size / self.sector_size)
@@ -35,10 +35,8 @@ def __init__(self, size: int) -> None:
         self.primary_array_byte = int(self.primary_array_lba * self.sector_size)
         self.backup_header_lba = int(self.total_lba - 1)
         self.backup_header_byte = int(self.backup_header_lba * self.sector_size)
-        self.backup_header_array_lba = int(self.total_lba - 33)
-        self.backup_header_array_byte = int(
-            self.backup_header_array_lba * self.sector_size
-        )
+        self.backup_array_lba = int(self.total_lba - 33)
+        self.backup_array_byte = int(self.backup_array_lba * self.sector_size)
 
 
 class Disk:
@@ -53,7 +51,9 @@ class Disk:
         geometry:
     """
 
-    def __init__(self, image_path: str, size: int = 0) -> None:
+    def __init__(
+        self, image_path: str, size: int = 0, fresh_disk: bool = False
+    ) -> None:
         """Init Disk with a file path and size in bytes"""
         # @TODO: check that disk is large enough to contain all table data
         self.image_path = pathlib.Path(image_path)
@@ -62,8 +62,6 @@ def __init__(self, image_path: str, size: int = 0) -> None:
         if self.image_path.exists():
             self.size = self.image_path.stat().st_size
         self.geometry = Geometry(self.size)
-        self.write()
-
-    def write(self) -> None:
-        """Write blank disk"""
-        self.image_path.write_bytes(b"\x00" * self.geometry.total_bytes)
+        # @NOTE: this is unnecessary
+        if not self.image_path.exists() or fresh_disk:
+            self.image_path.write_bytes(b"\x00" * self.geometry.total_bytes)

gpt_image/table.py

@@ -5,7 +5,7 @@
 import binascii
 import uuid
 from dataclasses import dataclass
-from typing import List, Tuple
+from typing import List
 
 from gpt_image.disk import Disk, Geometry
 
@@ -115,9 +115,9 @@ def __init__(self, geometry: Geometry, is_backup: bool = False):
                 self.secondary_header_lba.data,
                 self.primary_header_lba.data,
             )
-            self.partition_array_start.data = (
-                self.geometry.backup_header_array_lba
-            ).to_bytes(8, "little")
+            self.partition_array_start.data = (self.geometry.backup_array_lba).to_bytes(
+                8, "little"
+            )
 
         # header start byte relative the table itself, not the disk
         # primary will be 0 secondary will be LBA 32
@@ -183,6 +183,8 @@ def __init__(
         )
         if not partition_guid:
             self.partition_guid = TableEntry(16, 16, uuid.uuid4().bytes_le)
+        else:
+            self.partition_guid = TableEntry(16, 16, partition_guid.bytes_le)
         self.alignment = alignment
         self._partition_list = partition_list
         self._partition_size = size
@@ -209,7 +211,7 @@ def _get_partition_lba(self) -> tuple:
         """Calculate the partition's LBAs"""
         last_partition = self._start_lba.to_bytes(4, "little")
         if self._partition_list:
-            last_partition = self.partition_list[-1].last_lba.data
+            last_partition = self._partition_list[-1].last_lba.data
         first_lba = int.from_bytes(last_partition, "little") + 1
         # calculate partition size in LBA with alignment considered
         total_lba = int(self._partition_size / (self._sector_size * self.alignment) + 1)
@@ -224,9 +226,9 @@ def as_bytes(self) -> bytes:
 class Table:
     """GPT Partition Table Object
 
-    A table contains a primary and secondary header and a primary
-    and secondary partition entry table.  All management of their entries
-    is done with this object.
+    The Table class the meant to be used by the consumer.  The underlying
+    classes should be called through functions in this class and not
+    directly used.
     """
 
     def __init__(self, disk: Disk, sector_size: int = 512) -> None:
@@ -267,7 +269,13 @@ def write(self):
             f.seek(self.geometry.backup_header_byte)
             f.write(self.secondary_header.as_bytes())
 
-    def create_partition(self, name: str, size: int, guid: str, alignment: int = 8):
+            # write secondary partition table
+            # f.seek(self.geometry.backup_array_byte)
+            # f.write(b"".join(self.partition_entries))
+
+    def create_partition(
+        self, name: str, size: int, guid: uuid.UUID, alignment: int = 8
+    ):
         part = Partition(
             self.partition_entries,
             name,

tests/test_integration.py

@@ -1,9 +1,11 @@
+import uuid
+
 from gpt_image import disk, table
 
 
 def test_e2e():
-    new_disk = disk.Disk("test-disk.raw", 2 * 1024 * 1024)
+    new_disk = disk.Disk("test-disk.raw", 2 * 1024 * 1024, True)
 
     t = table.Table(new_disk)
-    # t.create_partition("test-part", 2 * 1024, None)
+    # t.create_partition("test-part", 2 * 1024, uuid.uuid4())
     t.write()