Separate partition module

gpt_image/entry.py

@@ -0,0 +1,14 @@
+class Entry:
+    def __init__(self, offset: int, length: int, data: bytes):
+        self.offset = offset
+        self.length = length
+        self.data = data
+
+    def int_to_bytes(self, number: int) -> None:
+        """Convert in to the proper byte length"""
+        self.data = (number).to_bytes(self.length, "little")
+
+    def str_to_bytes(self, string: str) -> None:
+        b_string = bytes(string, encoding="utf_16_le")
+        padded = b_string + bytes(self.length - len(b_string))
+        self.data = padded

gpt_image/partition.py

@@ -1,10 +1,9 @@
 import uuid
 from dataclasses import dataclass
 from math import ceil
-from tkinter import W
 
 from gpt_image.disk import Geometry
-from gpt_image.table import TableEntry
+from gpt_image.entry import Entry
 
 
 class Partition:
@@ -41,12 +40,12 @@ def __init__(
           size: partition size in Bytes
         """
         # create an empty partition object
-        self.type_guid = TableEntry(0, 16, b"\x00" * 16)
-        self.partition_guid = TableEntry(16, 16, b"\x00" * 16)
-        self.first_lba = TableEntry(32, 8, b"\x00" * 8)
-        self.last_lba = TableEntry(40, 8, b"\x00" * 8)
-        self.attribute_flags = TableEntry(48, 8, b"\x00" * 8)
-        self.partition_name = TableEntry(56, 72, b"\x00" * 72)
+        self.type_guid = Entry(0, 16, b"\x00" * 16)
+        self.partition_guid = Entry(16, 16, b"\x00" * 16)
+        self.first_lba = Entry(32, 8, b"\x00" * 8)
+        self.last_lba = Entry(40, 8, b"\x00" * 8)
+        self.attribute_flags = Entry(48, 8, b"\x00" * 8)
+        self.partition_name = Entry(56, 72, b"\x00" * 72)
 
         # if name is set, this isn't an empty partition. Set relevant fields
         if name:
@@ -110,7 +109,7 @@ def _get_first_lba(self, partition: Partition) -> int:
 
         def next_lba(end_lba: int, alignment: int):
             m = int(end_lba / alignment)
-            return m * alignment
+            return (m + 1) * alignment
 
         largest_lba = 0
         for part in self.entries:
@@ -132,7 +131,8 @@ def _get_last_lba(self, partition: Partition) -> int:
 
         # round the LBA up to ensure our LBA will hold the partition
         lba = ceil(partition.size / self._geometry.sector_size)
-        return (partition.first_lba + lba) - 1
+        f_lba = int.from_bytes(partition.first_lba.data, byteorder="little")
+        return (f_lba + lba) - 1
 
     def _get_next_partition(self) -> int:
         """Return the index of the next unused partition"""

gpt_image/table.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 
 from gpt_image.disk import Disk, Geometry
+from gpt_image.partition import Partition, PartitionEntry
 
 
 @dataclass
@@ -153,73 +154,6 @@ def as_bytes(self) -> bytes:
         return b"".join(byte_list)
 
 
-class Partition:
-    """Partition class represents a GPT partition
-
-    Start and end LBA are set to None because they must be calculated
-    from a table's partition list.
-    """
-
-    @dataclass
-    class Type:
-        """GPT Partition Types
-
-        https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_entries
-        """
-
-        LinuxFileSystem = "0FC63DAF-8483-4772-8E79-3D69D8477DE4"
-        EFISystemPartition = "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
-
-    def __init__(
-        self,
-        name: str = None,
-        size: int = 0,
-        partition_guid: uuid.UUID = None,
-        alignment: int = 8,
-    ):
-        """Initialize Partition Object
-
-        All parameters have a default value to allow Partition() to create
-        an empty partition object.  If "name" is set, we assume this is not
-        an empty object and set the other values.
-        """
-        # create an empty partition object
-        self.type_guid = TableEntry(0, 16, b"\x00" * 16)
-        self.partition_guid = TableEntry(16, 16, b"\x00" * 16)
-        self.first_lba = TableEntry(32, 8, b"\x00" * 8)
-        self.last_lba = TableEntry(40, 8, b"\x00" * 8)
-        self.attribute_flags = TableEntry(48, 8, b"\x00" * 8)
-        self.partition_name = TableEntry(56, 72, b"\x00" * 72)
-
-        # if name is set, this isn't an empty partition. Set relevant fields
-        if name:
-            self.type_guid.data = uuid.UUID(Partition.Type.LinuxFileSystem).bytes_le
-            if not partition_guid:
-                self.partition_guid.data = uuid.uuid4().bytes_le
-            else:
-                self.partition_guid.data = partition_guid.bytes_le
-            b_name = bytes(name, encoding="utf_16_le")
-            # ensure the partition name is padded
-            self.partition_name.data = b_name + bytes(72 - len(b_name))
-
-        self.alignment = alignment
-        self.size = size
-
-        self.partition_fields = [
-            self.type_guid,
-            self.partition_guid,
-            self.first_lba,
-            self.last_lba,
-            self.attribute_flags,
-            self.partition_name,
-        ]
-
-    def as_bytes(self) -> bytes:
-        """Return the partition as bytes"""
-        byte_list = [x.data for x in self.partition_fields]
-        return b"".join(byte_list)
-
-
 class Table:
     """GPT Partition Table Object
 
@@ -235,7 +169,7 @@ def __init__(self, disk: Disk, sector_size: int = 512) -> None:
         self.primary_header = Header(self.geometry)
         self.secondary_header = Header(self.geometry, is_backup=True)
 
-        self.partition_entries = [Partition()] * 128
+        self.partitions = PartitionEntry(self.geometry)
 
     def write(self):
         """Write the table to disk"""
@@ -260,15 +194,15 @@ def write(self):
 
             # write primary partition table
             f.seek(self.geometry.primary_array_byte)
-            f.write(self._partition_entries_as_bytes())
+            f.write(self.partitions.as_bytes())
 
             # move to secondary header location and write
             f.seek(self.geometry.backup_header_byte)
             f.write(self.secondary_header.as_bytes())
 
             # write secondary partition table
             f.seek(self.geometry.backup_array_byte)
-            f.write(self._partition_entries_as_bytes())
+            f.write(self.partitions.as_bytes())
 
     def create_partition(
         self, name: str, size: int, guid: uuid.UUID, alignment: int = 8
@@ -279,50 +213,11 @@ def create_partition(
             guid,
             alignment,
         )
-        # find the first empty partition index
-        for idx, partition in enumerate(self.partition_entries):
-            if int.from_bytes(partition.partition_name.data, byteorder="little") == 0:
-                part.first_lba.data = self._first_lba()
-                last_lba = int(part.size / self.geometry.sector_size) + int.from_bytes(
-                    part.first_lba.data, byteorder="little"
-                )
-                # @NOTE: last_lba must consider alignment
-                part.last_lba.data = (last_lba).to_bytes(8, "little")
-                self.partition_entries[idx] = part
-                break
-
-    def _first_lba(self) -> bytes:
-        """Calculate the first LBA of a new partition
-
-        Search for the largest LBA, this will be used to calculate the first
-        LBA of the partition being created.  If it is 0, all partitions are empty
-        and we must start at LBA 34.
-
-        Example: the partition with the largest LBA is LBA 412. The new partition's
-        LBA will be set to 413 (412 + 1)
-
-        """
-        last_lba = 0
-        for partition in self.partition_entries:
-            last_lba_int = int.from_bytes(partition.last_lba.data, byteorder="little")
-            if last_lba_int > last_lba:
-                last_lba = last_lba_int
-        if last_lba == 0:
-            return (34).to_bytes(8, "little")
-        # @NOTE: this is NOT proper alignment
-        return (last_lba + 1).to_bytes(8, "little")
-
-    def _last_lba(self, size: int) -> bytes:
-        """Calculate the last LBA of a new partition"""
-        pass
-
-    def _partition_entries_as_bytes(self):
-        parts = [x.as_bytes() for x in self.partition_entries]
-        return b"".join(parts)
+        self.partitions.add(part)
 
     def checksum_partitions(self, header: Header):
         """Checksum the partition entries"""
-        part_entry_bytes = self._partition_entries_as_bytes()
+        part_entry_bytes = self.partitions.as_bytes()
         header.partition_array_crc.data = binascii.crc32(part_entry_bytes).to_bytes(
             4, "little"
         )

tests/test_disk.py

@@ -30,7 +30,6 @@ def test_disk_new(tmp_path):
     image_name = tmp_path / "test.img"
     abs_path = image_name.resolve()
     disk = Disk(str(abs_path), image_size)
-    assert Path(disk.name).exists()
     assert disk.size == image_size
     assert disk.image_path == image_name
     assert image_name.stat().st_size == image_size