Some QCOW2 improvements and unit tests

dissect/hypervisor/disk/c_qcow2.py

@@ -171,9 +171,6 @@
 )
 
 
-def ctz(value: int, size: int = 32) -> int:
+def ctz(value: int) -> int:
     """Count the number of zero bits in an integer of a given size."""
-    for i in range(size):
-        if value & (1 << i):
-            return i
-    return 0
+    return (value & -value).bit_length() - 1 if value else 0

dissect/hypervisor/disk/c_qcow2.pyi

@@ -0,0 +1,190 @@
+# Generated by cstruct-stubgen
+from typing import BinaryIO, Literal, overload
+
+import dissect.cstruct as __cs__
+from typing_extensions import TypeAlias
+
+class _c_qcow2(__cs__.cstruct):
+    MIN_CLUSTER_BITS: Literal[9] = ...
+    MAX_CLUSTER_BITS: Literal[21] = ...
+    QCOW2_COMPRESSED_SECTOR_SIZE: Literal[512] = ...
+    QCOW2_COMPRESSION_TYPE_ZLIB: Literal[0] = ...
+    QCOW2_COMPRESSION_TYPE_ZSTD: Literal[1] = ...
+    L1E_SIZE: Literal[8] = ...
+    L2E_SIZE_NORMAL: Literal[8] = ...
+    L2E_SIZE_EXTENDED: Literal[16] = ...
+    L1E_OFFSET_MASK: Literal[72057594037927424] = ...
+    L2E_OFFSET_MASK: Literal[72057594037927424] = ...
+    L2E_COMPRESSED_OFFSET_SIZE_MASK: Literal[4611686018427387903] = ...
+    QCOW_OFLAG_COPIED: Literal[9223372036854775808] = ...
+    QCOW_OFLAG_COMPRESSED: Literal[4611686018427387904] = ...
+    QCOW_OFLAG_ZERO: Literal[1] = ...
+    QCOW_EXTL2_SUBCLUSTERS_PER_CLUSTER: Literal[32] = ...
+    QCOW2_INCOMPAT_DIRTY_BITNR: Literal[0] = ...
+    QCOW2_INCOMPAT_CORRUPT_BITNR: Literal[1] = ...
+    QCOW2_INCOMPAT_DATA_FILE_BITNR: Literal[2] = ...
+    QCOW2_INCOMPAT_COMPRESSION_BITNR: Literal[3] = ...
+    QCOW2_INCOMPAT_EXTL2_BITNR: Literal[4] = ...
+    QCOW2_INCOMPAT_DIRTY: Literal[1] = ...
+    QCOW2_INCOMPAT_CORRUPT: Literal[2] = ...
+    QCOW2_INCOMPAT_DATA_FILE: Literal[4] = ...
+    QCOW2_INCOMPAT_COMPRESSION: Literal[8] = ...
+    QCOW2_INCOMPAT_EXTL2: Literal[16] = ...
+    QCOW2_EXT_MAGIC_END: Literal[0] = ...
+    QCOW2_EXT_MAGIC_BACKING_FORMAT: Literal[3799591626] = ...
+    QCOW2_EXT_MAGIC_FEATURE_TABLE: Literal[1745090647] = ...
+    QCOW2_EXT_MAGIC_CRYPTO_HEADER: Literal[87539319] = ...
+    QCOW2_EXT_MAGIC_BITMAPS: Literal[595929205] = ...
+    QCOW2_EXT_MAGIC_DATA_FILE: Literal[1145132097] = ...
+    class QCowHeader(__cs__.Structure):
+        magic: _c_qcow2.uint32
+        version: _c_qcow2.uint32
+        backing_file_offset: _c_qcow2.uint64
+        backing_file_size: _c_qcow2.uint32
+        cluster_bits: _c_qcow2.uint32
+        size: _c_qcow2.uint64
+        crypt_method: _c_qcow2.uint32
+        l1_size: _c_qcow2.uint32
+        l1_table_offset: _c_qcow2.uint64
+        refcount_table_offset: _c_qcow2.uint64
+        refcount_table_clusters: _c_qcow2.uint32
+        nb_snapshots: _c_qcow2.uint32
+        snapshots_offset: _c_qcow2.uint64
+        incompatible_features: _c_qcow2.uint64
+        compatible_features: _c_qcow2.uint64
+        autoclear_features: _c_qcow2.uint64
+        refcount_order: _c_qcow2.uint32
+        header_length: _c_qcow2.uint32
+        compression_type: _c_qcow2.uint8
+        padding: __cs__.Array[_c_qcow2.uint8]
+        @overload
+        def __init__(
+            self,
+            magic: _c_qcow2.uint32 | None = ...,
+            version: _c_qcow2.uint32 | None = ...,
+            backing_file_offset: _c_qcow2.uint64 | None = ...,
+            backing_file_size: _c_qcow2.uint32 | None = ...,
+            cluster_bits: _c_qcow2.uint32 | None = ...,
+            size: _c_qcow2.uint64 | None = ...,
+            crypt_method: _c_qcow2.uint32 | None = ...,
+            l1_size: _c_qcow2.uint32 | None = ...,
+            l1_table_offset: _c_qcow2.uint64 | None = ...,
+            refcount_table_offset: _c_qcow2.uint64 | None = ...,
+            refcount_table_clusters: _c_qcow2.uint32 | None = ...,
+            nb_snapshots: _c_qcow2.uint32 | None = ...,
+            snapshots_offset: _c_qcow2.uint64 | None = ...,
+            incompatible_features: _c_qcow2.uint64 | None = ...,
+            compatible_features: _c_qcow2.uint64 | None = ...,
+            autoclear_features: _c_qcow2.uint64 | None = ...,
+            refcount_order: _c_qcow2.uint32 | None = ...,
+            header_length: _c_qcow2.uint32 | None = ...,
+            compression_type: _c_qcow2.uint8 | None = ...,
+            padding: __cs__.Array[_c_qcow2.uint8] | None = ...,
+        ): ...
+        @overload
+        def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+    class QCowExtension(__cs__.Structure):
+        magic: _c_qcow2.uint32
+        len: _c_qcow2.uint32
+        @overload
+        def __init__(self, magic: _c_qcow2.uint32 | None = ..., len: _c_qcow2.uint32 | None = ...): ...
+        @overload
+        def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+    class Qcow2CryptoHeaderExtension(__cs__.Structure):
+        offset: _c_qcow2.uint64
+        length: _c_qcow2.uint64
+        @overload
+        def __init__(self, offset: _c_qcow2.uint64 | None = ..., length: _c_qcow2.uint64 | None = ...): ...
+        @overload
+        def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+    class Qcow2BitmapHeaderExt(__cs__.Structure):
+        nb_bitmaps: _c_qcow2.uint32
+        reserved32: _c_qcow2.uint32
+        bitmap_directory_size: _c_qcow2.uint64
+        bitmap_directory_offset: _c_qcow2.uint64
+        @overload
+        def __init__(
+            self,
+            nb_bitmaps: _c_qcow2.uint32 | None = ...,
+            reserved32: _c_qcow2.uint32 | None = ...,
+            bitmap_directory_size: _c_qcow2.uint64 | None = ...,
+            bitmap_directory_offset: _c_qcow2.uint64 | None = ...,
+        ): ...
+        @overload
+        def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+    class QCowSnapshotHeader(__cs__.Structure):
+        l1_table_offset: _c_qcow2.uint64
+        l1_size: _c_qcow2.uint32
+        id_str_size: _c_qcow2.uint16
+        name_size: _c_qcow2.uint16
+        date_sec: _c_qcow2.uint32
+        date_nsec: _c_qcow2.uint32
+        vm_clock_nsec: _c_qcow2.uint64
+        vm_state_size: _c_qcow2.uint32
+        extra_data_size: _c_qcow2.uint32
+        @overload
+        def __init__(
+            self,
+            l1_table_offset: _c_qcow2.uint64 | None = ...,
+            l1_size: _c_qcow2.uint32 | None = ...,
+            id_str_size: _c_qcow2.uint16 | None = ...,
+            name_size: _c_qcow2.uint16 | None = ...,
+            date_sec: _c_qcow2.uint32 | None = ...,
+            date_nsec: _c_qcow2.uint32 | None = ...,
+            vm_clock_nsec: _c_qcow2.uint64 | None = ...,
+            vm_state_size: _c_qcow2.uint32 | None = ...,
+            extra_data_size: _c_qcow2.uint32 | None = ...,
+        ): ...
+        @overload
+        def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+    class QCowSnapshotExtraData(__cs__.Structure):
+        vm_state_size_large: _c_qcow2.uint64
+        disk_size: _c_qcow2.uint64
+        icount: _c_qcow2.uint64
+        @overload
+        def __init__(
+            self,
+            vm_state_size_large: _c_qcow2.uint64 | None = ...,
+            disk_size: _c_qcow2.uint64 | None = ...,
+            icount: _c_qcow2.uint64 | None = ...,
+        ): ...
+        @overload
+        def __init__(self, fh: bytes | memoryview | bytearray | BinaryIO, /): ...
+
+    class QCow2ClusterType(__cs__.Enum):
+        QCOW2_CLUSTER_UNALLOCATED = ...
+        QCOW2_CLUSTER_ZERO_PLAIN = ...
+        QCOW2_CLUSTER_ZERO_ALLOC = ...
+        QCOW2_CLUSTER_NORMAL = ...
+        QCOW2_CLUSTER_COMPRESSED = ...
+
+    class QCow2SubclusterType(__cs__.Enum):
+        QCOW2_SUBCLUSTER_UNALLOCATED_PLAIN = ...
+        QCOW2_SUBCLUSTER_UNALLOCATED_ALLOC = ...
+        QCOW2_SUBCLUSTER_ZERO_PLAIN = ...
+        QCOW2_SUBCLUSTER_ZERO_ALLOC = ...
+        QCOW2_SUBCLUSTER_NORMAL = ...
+        QCOW2_SUBCLUSTER_COMPRESSED = ...
+        QCOW2_SUBCLUSTER_INVALID = ...
+
+# Technically `c_qcow2` is an instance of `_c_qcow2`, but then we can't use it in type hints
+c_qcow2: TypeAlias = _c_qcow2
+
+QCOW2_MAGIC: int
+QCOW2_MAGIC_BYTES: bytes
+
+QCOW2_INCOMPAT_MASK: int
+
+QCow2ClusterType: TypeAlias = _c_qcow2.QCow2ClusterType
+QCow2SubclusterType: TypeAlias = _c_qcow2.QCow2SubclusterType
+
+NORMAL_SUBCLUSTER_TYPES: tuple[QCow2SubclusterType, ...]
+ZERO_SUBCLUSTER_TYPES: tuple[QCow2SubclusterType, ...]
+UNALLOCATED_SUBCLUSTER_TYPES: tuple[QCow2SubclusterType, ...]
+
+def ctz(value: int) -> int: ...