Add endian keyword argument to type reads

Author: Schamper

dissect/cstruct/bitbuffer.py

@@ -9,9 +9,10 @@
 class BitBuffer:
     """Implements a bit buffer that can read and write bit fields."""
 
-    def __init__(self, stream: BinaryIO, endian: str):
+    def __init__(self, stream: BinaryIO, *, endian: str, **kwargs):
         self.stream = stream
         self.endian = endian
+        self.kwargs = kwargs
 
         self._type: type[BaseType] | None = None
         self._buffer = 0
@@ -24,7 +25,7 @@ def read(self, field_type: type[BaseType], bits: int) -> int:
 
             self._type = field_type
             self._remaining = field_type.size * 8
-            self._buffer = field_type._read(self.stream)
+            self._buffer = field_type._read(self.stream, endian=self.endian, **self.kwargs)
 
         if isinstance(self._buffer, bytes):
             if self.endian == "<":
@@ -71,7 +72,7 @@ def write(self, field_type: type[BaseType], data: int, bits: int) -> None:
 
     def flush(self) -> None:
         if self._type is not None:
-            self._type._write(self.stream, self._buffer)
+            self._type._write(self.stream, self._buffer, endian=self.endian, **self.kwargs)
         self._type = None
         self._remaining = 0
         self._buffer = 0

dissect/cstruct/compiler.py

@@ -119,7 +119,7 @@ def generate_source(self) -> str:
         """
 
         if any(field.bits for field in self.fields):
-            preamble += "bit_reader = BitBuffer(stream, cls.cs.endian)\n"
+            preamble += "bit_reader = BitBuffer(stream, endian=endian, **kwargs)\n"
 
         read_code = "\n".join(self._generate_fields())
 
@@ -132,7 +132,7 @@ def generate_source(self) -> str:
 
         code = indent(dedent(preamble).lstrip() + read_code + dedent(outro), "    ")
 
-        return f"def _read(cls, stream, context=None):\n{code}"
+        return f"def _read(cls, stream, *, context=None, endian, **kwargs):\n{code}"
 
     def _generate_fields(self) -> Iterator[str]:
         current_offset = 0
@@ -229,7 +229,7 @@ def align_to_field(field: Field) -> Iterator[str]:
     def _generate_structure(self, field: Field) -> Iterator[str]:
         template = f"""
         {"_s = stream.tell()" if field.type.dynamic else ""}
-        r["{field._name}"] = {self._map_field(field)}._read(stream, context=r)
+        r["{field._name}"] = {self._map_field(field)}._read(stream, context=r, endian=endian, **kwargs)
         {f's["{field._name}"] = stream.tell() - _s' if field.type.dynamic else ""}
         """
 
@@ -238,7 +238,7 @@ def _generate_structure(self, field: Field) -> Iterator[str]:
     def _generate_array(self, field: Field) -> Iterator[str]:
         template = f"""
         {"_s = stream.tell()" if field.type.dynamic else ""}
-        r["{field._name}"] = {self._map_field(field)}._read(stream, context=r)
+        r["{field._name}"] = {self._map_field(field)}._read(stream, context=r, endian=endian, **kwargs)
         {f's["{field._name}"] = stream.tell() - _s' if field.type.dynamic else ""}
         """
 
@@ -311,7 +311,7 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
                     item_parser = parser_template.format(type="_et", getter=f"_b[i:i + {field_type.type.size}]")
                     list_comp = f"[{item_parser} for i in range(0, {count}, {field_type.type.size})]"
                 elif issubclass(field_type.type, Pointer):
-                    item_parser = "_et.__new__(_et, e, stream, r)"
+                    item_parser = "_et.__new__(_et, e, stream, context=r, endian=endian, **kwargs)"
                     list_comp = f"[{item_parser} for e in {getter}]"
                 else:
                     item_parser = parser_template.format(type="_et", getter="e")
@@ -322,7 +322,7 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
                 parser = f"type.__call__({self._map_field(field)}, {getter})"
             elif issubclass(field_type, Pointer):
                 reads.append(f"_pt = {self._map_field(field)}")
-                parser = f"_pt.__new__(_pt, {getter}, stream, r)"
+                parser = f"_pt.__new__(_pt, {getter}, stream, context=r, endian=endian, **kwargs)"
             else:
                 parser = parser_template.format(type=self._map_field(field), getter=getter)
 
@@ -335,7 +335,7 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
         if fmt == "x" or (len(fmt) == 2 and fmt[1] == "x"):
             unpack = ""
         else:
-            unpack = f'data = _struct(cls.cs.endian, "{fmt}").unpack(buf)\n'
+            unpack = f'data = _struct(endian, "{fmt}").unpack(buf)\n'
 
         template = f"""
         buf = stream.read({size})

dissect/cstruct/cstruct.py

@@ -2,13 +2,16 @@
 
 import copy
 import ctypes as _ctypes
+import inspect
 import struct
 import sys
 import types
+import warnings
+from contextlib import contextmanager
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, BinaryIO, TypeVar, cast
+from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypeVar, cast
 
-from dissect.cstruct.exceptions import ResolveError
+from dissect.cstruct.exceptions import Error, ResolveError
 from dissect.cstruct.expression import Expression
 from dissect.cstruct.parser import CStyleParser
 from dissect.cstruct.types import (
@@ -26,10 +29,10 @@
     Void,
     Wchar,
 )
-from dissect.cstruct.types.base import MetaType
+from dissect.cstruct.types.base import MetaType, normalize_endianness
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
     from typing import TypeAlias
 
     from dissect.cstruct.types import (
@@ -40,20 +43,23 @@
 
 T = TypeVar("T", bound=BaseType)
 
+AllowedEndianness: TypeAlias = Literal["little", "big", "network", "<", ">", "!", "@", "="]
+Endianness: TypeAlias = Literal["<", ">", "!", "@", "="]
+
 
 class cstruct:
     """Main class of cstruct. All types are registered in here.
 
     Args:
-        endian: The endianness to use when parsing.
+        endian: The endianness to use when parsing (little, big, network, <, >, !, @ or =).
         pointer: The pointer type to use for pointers.
     """
 
     DEF_CSTYLE = 1
     DEF_LEGACY = 2
 
-    def __init__(self, load: str = "", *, endian: str = "<", pointer: str | None = None):
-        self.endian = endian
+    def __init__(self, load: str = "", *, endian: AllowedEndianness = "<", pointer: str | None = None):
+        self.endian = normalize_endianness(endian)
 
         self.consts = {}
         self.includes = []
@@ -215,9 +221,11 @@ def __copy__(self) -> cstruct:
 
         # Update typedefs to point to the new cstruct instance
         for name, type in self.typedefs.items():
+            # Skip all the default types
             if name in cs.typedefs:
                 continue
 
+            # Update the cstruct instance for all other types
             if isinstance(type, MetaType):
                 new_type = copy.copy(type)
                 new_type.cs = cs
@@ -264,6 +272,33 @@ def add_custom_type(
             alignment: The alignment of the type.
             **kwargs: Additional attributes to add to the type.
         """
+        # In cstruct 5.0 we changed the function signature of _read and _write
+        # Check if the function signature is compatible, and throw an error if not
+        for type_to_check in (type_, type_.ArrayType):
+            type_name = type_.__name__ + (f".{type_.ArrayType.__name__}" if type_to_check is type_.ArrayType else "")
+
+            for method in ("_read", "_read_array", "_read_0", "_write", "_write_array", "_write_0"):
+                if not hasattr(type_to_check, method):
+                    continue
+
+                signature = inspect.signature(getattr(type_to_check, method))
+
+                # We added a few keyword-only parameters to the function signature, but any custom type will
+                # continue to work fine as long as they accept **kwargs
+                if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()):
+                    raise Error(
+                        f"Custom type {type_name} has an incompatible {method} method signature. "
+                        "Please refer to the changelog of dissect.cstruct 5.0 for more information."
+                    )
+
+                # Only warn if the method doesn't accept an endian parameter
+                if "endian" not in signature.parameters:
+                    warnings.warn(
+                        f"Custom type {type_name} is missing the 'endian' keyword-only parameter in its {method} method. "  # noqa: E501
+                        "Please refer to the changelog of dissect.cstruct 5.0 for more information.",
+                        stacklevel=2,
+                    )
+
         self.add_type(name, self._make_type(name, (type_,), size, alignment=alignment, attrs=kwargs))
 
     def load(self, definition: str, deftype: int | None = None, **kwargs) -> cstruct:
@@ -351,6 +386,25 @@ def copy(self) -> cstruct:
         """
         return copy.copy(self)
 
+    @contextmanager
+    def endianness(self, endian: AllowedEndianness) -> Iterator[cstruct]:
+        """Context manager for temporarily changing the endianness of this cstruct instance."""
+        new = self.copy()
+        new.endian = normalize_endianness(endian)
+        yield new
+
+    @contextmanager
+    def big_endian(self) -> Iterator[cstruct]:
+        """Context manager for temporarily changing the endianness to big endian."""
+        with self.endianness(">") as cs:
+            yield cs
+
+    @contextmanager
+    def little_endian(self) -> Iterator[cstruct]:
+        """Context manager for temporarily changing the endianness to little endian."""
+        with self.endianness("<") as cs:
+            yield cs
+
     def _make_type(
         self,
         name: str,