Add support for auto-detecting 4k GPT disks

dissect/volume/disk/disk.py

@@ -10,29 +10,72 @@
 
 
 class Disk:
-    def __init__(self, fh: BinaryIO, sector_size: int = 512):
+    """Generic disk partitioning implementation. The partition scheme is detected automatically.
+
+    Supported partition schemes:
+        - MBR (Master Boot Record)
+        - GPT (GUID Partition Table)
+        - APM (Apple Partition Map)
+        - BSD disklabel (contained in another partition or standalone)
+
+    Args:
+        fh: File-like object of a disk containing a partition scheme.
+        sector_size: Sector size in bytes.
+                     If not provided and the disk contains GPT, it will be detected automatically.
+                     Otherwise, 512 bytes is assumed.
+    """
+
+    def __init__(self, fh: BinaryIO, sector_size: int | None = None):
         self.fh = fh
-        self.sector_size = sector_size
+        self.sector_size = sector_size or 512
         self.scheme: APM | GPT | MBR | BSD = None
         self.partitions: list[Partition] = []
 
-        start = fh.tell()
         errors = []
 
-        # The GPT scheme also parses the protective MBR, so it must be tried before MBR.
-        # BSD is usually contained in another scheme's partition, but it can also live standalone.
-        # We try to detect BSD as part of another scheme later on, so only try to detect BSD last
-        # as standalone.
-        for scheme in [GPT, MBR, APM, BSD]:
-            try:
-                fh.seek(start)
-                self.scheme = scheme(fh, sector_size=self.sector_size)
+        # Do a little dance with MBR and GPT first, since we can try to determine the sector size here
+        try:
+            self.fh.seek(0)
+            self.scheme = MBR(self.fh, sector_size=self.sector_size)
+        except Exception as e:
+            errors.append(str(e))
 
-                break
-            except Exception as e:
-                errors.append(str(e))
+        if self.scheme and any(p.type == 0xEE for p in self.scheme.partitions):
+            # There's a protective MBR, potentially GPT
+            # Try to detect sector size until we find a valid GPT header
+            # If the user provided a sector size, we only try that one, otherwise we try the most common ones
+            for guess in [512, 4096] if sector_size is None else [sector_size]:
+                try:
+                    self.fh.seek(0)
+                    self.scheme = GPT(self.fh, sector_size=guess)
+                    # Winner winner chicken dinner
+                    self.sector_size = guess
+                    break
+                except Exception as e:
+                    errors.append(str(e))
+            else:
+                # No valid GPT found
+                if sector_size is None:
+                    reason = "Maybe exotic sector size?"
+                else:
+                    reason = f"Given sector size ({sector_size}) seems incorrect."
 
-        if not self.scheme:
+                raise DiskError(f"Found GPT type partition, but MBR scheme detected. {reason}")
+
+        else:
+            # It's not MBR or GPT, try the other schemes
+            # BSD is usually contained in another scheme's partition, but it can also live standalone.
+            # We try to detect BSD as part of another scheme later on, so only try to detect BSD last
+            # as standalone.
+            for scheme in [APM, BSD]:
+                try:
+                    self.fh.seek(0)
+                    self.scheme = scheme(self.fh, sector_size=self.sector_size)
+                    break
+                except Exception as e:
+                    errors.append(str(e))
+
+        if self.scheme is None:
             raise DiskError("Unable to detect a valid partition scheme:\n- {}".format("\n- ".join(errors)))
 
         main_scheme = self.scheme
@@ -44,9 +87,6 @@ def __init__(self, fh: BinaryIO, sector_size: int = 512):
             else:
                 self.partitions.append(partition)
 
-        if isinstance(self.scheme, MBR) and any(p.type == 0xEE for p in self.partitions):
-            raise DiskError("Found GPT type partition, but MBR scheme detected. Maybe 4K sector size.")
-
     @property
     def serial(self) -> int | None:
         if isinstance(self.scheme, MBR):

tests/disk/test_gpt.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import io
-import re
 from typing import BinaryIO
 from uuid import UUID
 
@@ -97,36 +96,35 @@ def test_hybrid_gpt(gpt_hybrid: BinaryIO) -> None:
 
 def test_gpt_4k(gpt_4k: BinaryIO) -> None:
     with pytest.raises(
-        disk.DiskError, match=re.escape("Found GPT type partition, but MBR scheme detected. Maybe 4K sector size.")
+        disk.DiskError,
+        match=r"Found GPT type partition, but MBR scheme detected. Given sector size \(512\) seems incorrect.",
     ):
-        disk.Disk(gpt_4k)
-
-    gpt_4k.seek(0)
-    d = disk.Disk(gpt_4k, sector_size=4096)
-
-    assert isinstance(d.scheme, GPT)
-    assert len(d.partitions) == 3
-
-    assert d.partitions[0].number == 1
-    assert d.partitions[0].offset == 0x100000
-    assert d.partitions[0].size == 0x100000
-    assert d.partitions[0].type == UUID("0fc63daf-8483-4772-8e79-3d69d8477de4")
-    assert d.partitions[0].guid == UUID("21b90a6e-0918-4e72-aa1a-85f8ba8ef8cc")
-    assert d.partitions[0].name == "Linux filesystem"
-
-    assert d.partitions[1].number == 2
-    assert d.partitions[1].offset == 0x300000
-    assert d.partitions[1].size == 0x100000
-    assert d.partitions[1].type == UUID("0fc63daf-8483-4772-8e79-3d69d8477de4")
-    assert d.partitions[1].guid == UUID("c6f4ad42-4652-448d-89d7-7cfa7710abe7")
-    assert d.partitions[1].name == "Linux filesystem"
-
-    assert d.partitions[2].number == 3
-    assert d.partitions[2].offset == 0x500000
-    assert d.partitions[2].size == 0xB5A000
-    assert d.partitions[2].type == UUID("0fc63daf-8483-4772-8e79-3d69d8477de4")
-    assert d.partitions[2].guid == UUID("b7230707-dcaa-4483-823b-06f9b718ee55")
-    assert d.partitions[2].name == "Linux filesystem"
+        disk.Disk(gpt_4k, sector_size=512)
+
+    for d in [disk.Disk(gpt_4k), disk.Disk(gpt_4k, sector_size=4096)]:
+        assert isinstance(d.scheme, GPT)
+        assert len(d.partitions) == 3
+
+        assert d.partitions[0].number == 1
+        assert d.partitions[0].offset == 0x100000
+        assert d.partitions[0].size == 0x100000
+        assert d.partitions[0].type == UUID("0fc63daf-8483-4772-8e79-3d69d8477de4")
+        assert d.partitions[0].guid == UUID("21b90a6e-0918-4e72-aa1a-85f8ba8ef8cc")
+        assert d.partitions[0].name == "Linux filesystem"
+
+        assert d.partitions[1].number == 2
+        assert d.partitions[1].offset == 0x300000
+        assert d.partitions[1].size == 0x100000
+        assert d.partitions[1].type == UUID("0fc63daf-8483-4772-8e79-3d69d8477de4")
+        assert d.partitions[1].guid == UUID("c6f4ad42-4652-448d-89d7-7cfa7710abe7")
+        assert d.partitions[1].name == "Linux filesystem"
+
+        assert d.partitions[2].number == 3
+        assert d.partitions[2].offset == 0x500000
+        assert d.partitions[2].size == 0xB5A000
+        assert d.partitions[2].type == UUID("0fc63daf-8483-4772-8e79-3d69d8477de4")
+        assert d.partitions[2].guid == UUID("b7230707-dcaa-4483-823b-06f9b718ee55")
+        assert d.partitions[2].name == "Linux filesystem"
 
 
 def test_gpt_esxi(gpt_esxi: BinaryIO) -> None: