Add endian keyword argument to type reads

Author: Schamper

dissect/cstruct/bitbuffer.py

@@ -9,7 +9,7 @@
 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):
         self.stream = stream
         self.endian = endian
 
@@ -24,7 +24,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)
 
         if isinstance(self._buffer, bytes):
             if self.endian == "<":
@@ -71,7 +71,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._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)\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):\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)
         {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)
         {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)"
                     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)"
             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,11 +2,14 @@
 
 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.exception import ResolveError
 from dissect.cstruct.expression import Expression
@@ -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 = []
@@ -216,9 +222,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
@@ -265,6 +273,25 @@ 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 warn 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))
+
+                # 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:
@@ -388,6 +415,28 @@ 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."""
+        original = self.endian
+        self.endian = normalize_endianness(endian)
+        try:
+            yield self
+        finally:
+            self.endian = original
+
+    @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,