From b6ee670da3b96b649c43132d5d47737bc0e5af02 Mon Sep 17 00:00:00 2001
From: Erik Schamper <1254028+Schamper@users.noreply.github.com>
Date: Tue, 4 Jun 2024 14:15:23 +0200
Subject: [PATCH] Compatibility with cstruct v4 (#28)

---
 dissect/xfs/c_xfs.py |  5 ++---
 dissect/xfs/xfs.py   | 22 +++++++++++++++-------
 pyproject.toml       | 10 ++++++++--
 tox.ini              |  1 +
 4 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/dissect/xfs/c_xfs.py b/dissect/xfs/c_xfs.py
index d0b2ddb..c84b972 100644
--- a/dissect/xfs/c_xfs.py
+++ b/dissect/xfs/c_xfs.py
@@ -1,6 +1,6 @@
 import stat
 
-from dissect import cstruct
+from dissect.cstruct import cstruct
 
 xfs_def = """
 typedef uint64 xfs_ino_t;                               /* <inode> type */
@@ -888,8 +888,7 @@
 } xfs_dir2_block_tail_t;
 """
 
-c_xfs = cstruct.cstruct(endian=">")
-c_xfs.load(xfs_def)
+c_xfs = cstruct(endian=">").load(xfs_def)
 
 FILETYPES = {
     0x0: None,
diff --git a/dissect/xfs/xfs.py b/dissect/xfs/xfs.py
index 7add956..d38fa58 100644
--- a/dissect/xfs/xfs.py
+++ b/dissect/xfs/xfs.py
@@ -9,7 +9,6 @@
 from typing import BinaryIO, Iterator, Optional, Union
 from uuid import UUID
 
-from dissect.cstruct import Instance
 from dissect.util import ts
 from dissect.util.stream import RangeStream, RunlistStream
 
@@ -95,7 +94,7 @@ def get_relative_inode(self, agnum: int, inum: int, *args, **kwargs) -> INode:
 
         return self.get_allocation_group(agnum).get_inode(inum, *args, **kwargs)
 
-    def walk_agi(self, block: int, agnum: int) -> Iterator[Instance]:
+    def walk_agi(self, block: int, agnum: int) -> Iterator[c_xfs.xfs_inobt_rec]:
         for record in self.walk_small_tree(block, agnum, 16, (c_xfs.XFS_IBT_MAGIC, c_xfs.XFS_IBT_CRC_MAGIC)):
             yield c_xfs.xfs_inobt_rec(record)
 
@@ -119,7 +118,11 @@ def walk_small_tree(
         yield from self._walk_small_tree(root, leaf_size, agnum, magic)
 
     def _walk_small_tree(
-        self, node: Instance, leaf_size: int, agnum: int, magic: Optional[list[int]] = None
+        self,
+        node: c_xfs.xfs_btree_sblock | c_xfs.xfs_btree_sblock_crc,
+        leaf_size: int,
+        agnum: int,
+        magic: Optional[list[int]] = None,
     ) -> Iterator[bytes]:
         fh = self.fh
         if magic and node.bb_magic not in magic:
@@ -141,7 +144,12 @@ def _walk_small_tree(
 
                 yield from self._walk_small_tree(self._sblock_s(fh), leaf_size, agnum, magic)
 
-    def _walk_large_tree(self, node: Instance, leaf_size: int, magic: Optional[list[int]] = None) -> Iterator[bytes]:
+    def _walk_large_tree(
+        self,
+        node: c_xfs.xfs_btree_lblock | c_xfs.xfs_btree_lblock_crc,
+        leaf_size: int,
+        magic: Optional[list[int]] = None,
+    ) -> Iterator[bytes]:
         fh = self.fh
         if magic and node.bb_magic not in magic:
             magic_values = ", ".join([f"0x{magic_value:x}" for magic_value in magic])
@@ -219,7 +227,7 @@ def walk_extents(self, fsb: int) -> Iterator[tuple[int, int, int, int]]:
         block = agnum * self.xfs.sb.sb_agblocks + blknum
         yield from self.xfs.walk_extents(block)
 
-    def walk_agi(self) -> Iterator[Instance]:
+    def walk_agi(self) -> Iterator[c_xfs.xfs_inobt_rec]:
         yield from self.xfs.walk_agi(self.agi.agi_root, self.num)
 
     def walk_tree(self, fsb: int, magic: Optional[list[int]] = None, small: bool = False):
@@ -260,7 +268,7 @@ def __init__(
     def __repr__(self) -> str:
         return f"<inode {self.inum} ({self.ag.num}:{self.relative_inum})>"
 
-    def _read_inode(self) -> Instance:
+    def _read_inode(self) -> c_xfs.xfs_dinode:
         self.ag.fh.seek(self.relative_inum * self.ag.sb.sb_inodesize)
         self._buf = io.BytesIO(self.ag.fh.read(self.ag.sb.sb_inodesize))
         inode = c_xfs.xfs_dinode(self._buf)
@@ -271,7 +279,7 @@ def _read_inode(self) -> Instance:
         return inode
 
     @property
-    def inode(self) -> Instance:
+    def inode(self) -> c_xfs.xfs_dinode:
         if not self._inode:
             self._inode = self._read_inode()
         return self._inode
diff --git a/pyproject.toml b/pyproject.toml
index 4f51964..12e98bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,8 +25,8 @@ classifiers = [
   "Topic :: Utilities",
 ]
 dependencies = [
-    "dissect.cstruct>=3.0.dev,<4.0.dev",
-    "dissect.util>=3.0.dev,<4.0.dev",
+    "dissect.cstruct>=4.dev,<5",
+    "dissect.util>=3,<4",
 ]
 dynamic = ["version"]
 
@@ -35,6 +35,12 @@ homepage = "https://dissect.tools"
 documentation = "https://docs.dissect.tools/en/latest/projects/dissect.xfs"
 repository = "https://github.com/fox-it/dissect.xfs"
 
+[project.optional-dependencies]
+dev = [
+    "dissect.cstruct>=4.0.dev,<5.0.dev",
+    "dissect.util>=3.0.dev,<4.0.dev",
+]
+
 [tool.black]
 line-length = 120
 
diff --git a/tox.ini b/tox.ini
index 67e8e8a..7bd2890 100644
--- a/tox.ini
+++ b/tox.ini
@@ -11,6 +11,7 @@ minversion = 4.4.3
 requires = virtualenv>=20.16.6
 
 [testenv]
+extras = dev
 deps =
     pytest
     pytest-cov