dissect.cstruct 4.8.dev9__py3-none-any.whl → 5.0.dev6__py3-none-any.whl

dissect/cstruct/__init__.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dissect.cstruct.bitbuffer import BitBuffer
 from dissect.cstruct.cstruct import cstruct, ctypes, ctypes_type
-from dissect.cstruct.exceptions import (
+from dissect.cstruct.exception import (
     Error,
     NullPointerDereference,
     ParserError,
@@ -28,7 +28,7 @@ from dissect.cstruct.types import (
     Wchar,
     WcharArray,
 )
-from dissect.cstruct.utils import (
+from dissect.cstruct.util import (
     dumpstruct,
     hexdump,
     p8,

dissect/cstruct/bitbuffer.py

@@ -9,7 +9,7 @@ if TYPE_CHECKING:
 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 @@ class BitBuffer:
 
             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 @@ class BitBuffer:
 
     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 @@ class _ReadSourceGenerator:
         """
 
         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 @@ class _ReadSourceGenerator:
 
         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 @@ class _ReadSourceGenerator:
     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 @@ class _ReadSourceGenerator:
     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 @@ class _ReadSourceGenerator:
                     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 @@ class _ReadSourceGenerator:
                 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 @@ class _ReadSourceGenerator:
         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,15 +2,18 @@ from __future__ import annotations
 
 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.exception import ResolveError
 from dissect.cstruct.expression import Expression
-from dissect.cstruct.parser import CStyleParser, TokenParser
+from dissect.cstruct.parser import CStyleParser
 from dissect.cstruct.types import (
     LEB128,
     BaseArray,
@@ -26,10 +29,10 @@ from dissect.cstruct.types import (
     Void,
     Wchar,
 )
-from dissect.cstruct.types.base import MetaType
+from dissect.cstruct.types.base import 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,26 +43,30 @@ if TYPE_CHECKING:
 
 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: dict[str, int | str | bytes] = {}
+        self.types: dict[str, type[BaseType]] = {}
+        self.includes: list[str] = []
 
-        self.consts = {}
-        self.lookups = {}
-        self.includes = []
         # fmt: off
-        self.typedefs = {
+        initial_types = {
             # Internal types
             "int8": self._make_packed_type("int8", "b", int),
             "uint8": self._make_packed_type("uint8", "B", int),
@@ -103,6 +110,22 @@ class cstruct:
             "signed long long": "int64",
             "unsigned long long": "uint64",
 
+            # Other convenience types
+            "u8": "uint8",
+            "u16": "uint16",
+            "u32": "uint32",
+            "u64": "uint64",
+            "u128": "uint128",
+            "__u8": "uint8",
+            "__u16": "uint16",
+            "__u32": "uint32",
+            "__u64": "uint64",
+            "__u128": "uint128",
+            "uchar": "uint8",
+            "ushort": "uint16",
+            "uint": "uint32",
+            "ulong": "uint32",
+
             # Windows types
             "BYTE": "uint8",
             "CHAR": "char",
@@ -170,24 +193,12 @@ class cstruct:
             "_DWORD": "uint32",
             "_QWORD": "uint64",
             "_OWORD": "uint128",
-
-            # Other convenience types
-            "u1": "uint8",
-            "u2": "uint16",
-            "u4": "uint32",
-            "u8": "uint64",
-            "u16": "uint128",
-            "__u8": "uint8",
-            "__u16": "uint16",
-            "__u32": "uint32",
-            "__u64": "uint64",
-            "uchar": "uint8",
-            "ushort": "uint16",
-            "uint": "uint32",
-            "ulong": "uint32",
         }
         # fmt: on
 
+        for name, type_ in initial_types.items():
+            self.add_type(name, type_)
+
         pointer = pointer or ("uint64" if sys.maxsize > 2**32 else "uint32")
         self.pointer: type[BaseType] = self.resolve(pointer)
         self._anonymous_count = 0
@@ -202,28 +213,25 @@ class cstruct:
             pass
 
         try:
-            return self.resolve(self.typedefs[attr])
+            return self.types[attr]
         except KeyError:
             pass
 
-        raise AttributeError(f"Invalid attribute: {attr}")
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute {attr!r}")
 
     def __copy__(self) -> cstruct:
         cs = cstruct(endian=self.endian, pointer=self.pointer.__name__)
         cs._anonymous_count = self._anonymous_count
-        cs.consts = self.consts.copy()
-        cs.lookups = self.lookups.copy()
         cs.includes = self.includes.copy()
 
-        # Update typedefs to point to the new cstruct instance
-        for name, type in self.typedefs.items():
-            if name in cs.typedefs:
-                continue
+        # Update types to point to the new cstruct instance
+        for name, type_ in self.types.items():
+            new_type = copy.copy(type_)
+            new_type.cs = cs
+            cs.add_type(name, new_type, replace=True)
 
-            if isinstance(type, MetaType):
-                new_type = copy.copy(type)
-                new_type.cs = cs
-                cs.typedefs[name] = new_type
+        for name, value in self.consts.items():
+            cs.add_const(name, value)
 
         return cs
 
@@ -232,22 +240,34 @@ class cstruct:
         self._anonymous_count += 1
         return name
 
+    def _add_attr(self, name: str, value: Any) -> None:
+        # Names that collide with the cstruct class (e.g. a struct named ``load``) are not set as attributes
+        # to avoid breaking the instance. They remain accessible through ``resolve`` and the type dicts.
+        if name in _RESERVED_NAMES:
+            return
+        setattr(self, name, value)
+
     def add_type(self, name: str, type_: type[BaseType] | str, replace: bool = False) -> None:
-        """Add a type or type reference.
+        """Add a type or type alias.
 
         Only use this method when creating type aliases or adding already bound types.
+        All types will be resolved to their actual type objects prior to being added.
 
         Args:
             name: Name of the type to be added.
             type_: The type to be added. Can be a str reference to another type or a compatible type class.
+                   If a str is given, it will be resolved to the actual type object.
+            replace: Whether to replace the type if it already exists.
 
         Raises:
             ValueError: If the type already exists.
         """
-        if not replace and (name in self.typedefs and self.resolve(self.typedefs[name]) != self.resolve(type_)):
+        typeobj = self.resolve(type_)
+        if not replace and (existing := self.types.get(name)) is not None and existing is not typeobj:
             raise ValueError(f"Duplicate type: {name}")
 
-        self.typedefs[name] = type_
+        self.types[name] = typeobj
+        self._add_attr(name, typeobj)
 
     addtype = add_type
 
@@ -266,19 +286,56 @@ class cstruct:
             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 add_const(self, name: str, value: Any) -> None:
+        """Add a constant value.
+
+        Args:
+            name: Name of the constant to be added.
+            value: The value of the constant.
+        """
+        self.consts[name] = value
+        self._add_attr(name, value)
+
+    def del_const(self, name: str) -> None:
+        """Delete a constant value.
+
+        Args:
+            name: Name of the constant to be deleted.
+
+        Raises:
+            KeyError: If the constant does not exist.
+        """
+        del self.consts[name]
+        self.__dict__.pop(name, None)
+
     def load(self, definition: str, deftype: int | None = None, **kwargs) -> cstruct:
         """Parse structures from the given definitions using the given definition type.
 
         Definitions can be parsed using different parsers. Currently, there's
         only one supported parser - DEF_CSTYLE. Parsers can add types and
         modify this cstruct instance. Arguments can be passed to parsers
-        using kwargs.
-
-        The CSTYLE parser was recently replaced with token based parser,
-        instead of a strictly regex based one. The old parser is still available
-        by using DEF_LEGACY.
+        using ``kwargs``.
 
         Args:
             definition: The definition to parse.
@@ -288,9 +345,9 @@ class cstruct:
         deftype = deftype or cstruct.DEF_CSTYLE
 
         if deftype == cstruct.DEF_CSTYLE:
-            TokenParser(self, **kwargs).parse(definition)
-        elif deftype == cstruct.DEF_LEGACY:
             CStyleParser(self, **kwargs).parse(definition)
+        else:
+            raise ValueError(f"Unknown definition type: {deftype}")
 
         return self
 
@@ -307,6 +364,42 @@ class cstruct:
         with Path(path).open() as fh:
             self.load(fh.read(), deftype, **kwargs)
 
+    def cdef(self) -> str:
+        """Render all constants, structure, union, enum and flag definitions back to their C-style definitions.
+
+        Note that the result is semantically equivalent to the original definitions, but not necessarily identical.
+
+        Returns:
+            The C-style definitions as a string.
+        """
+        empty = cstruct()
+
+        blocks = []
+
+        defines = []
+        for name, value in self.consts.items():
+            if name in empty.consts:
+                continue
+
+            defines.append(f"#define {name} {value!r}")
+
+        if defines:
+            blocks.append("\n".join(defines))
+
+        for name, typedef in self.types.items():
+            if name in empty.types or not isinstance(typedef, type):
+                continue
+
+            if not issubclass(typedef, (Structure, Enum, Flag)):
+                continue
+
+            if typedef.__name__ == name:
+                blocks.append(typedef.cdef())
+            else:
+                blocks.append(f"typedef {typedef.__name__} {name};")
+
+        return "\n\n".join(blocks)
+
     def read(self, name: str, stream: BinaryIO) -> Any:
         """Parse data using a given type.
 
@@ -334,20 +427,13 @@ class cstruct:
         Raises:
             ResolveError: If the type can't be resolved.
         """
-        type_name = name
-        if not isinstance(type_name, str):
-            return type_name
+        if not isinstance(name, str):
+            return name
 
-        for _ in range(10):
-            if type_name not in self.typedefs:
-                raise ResolveError(f"Unknown type {name}")
-
-            type_name = self.typedefs[type_name]
-
-            if not isinstance(type_name, str):
-                return type_name
-
-        raise ResolveError(f"Recursion limit exceeded while resolving type {name}")
+        try:
+            return self.types[name]
+        except KeyError:
+            raise ResolveError(f"Unknown type {name}") from None
 
     def copy(self) -> cstruct:
         """Create a copy of this cstruct instance.
@@ -357,6 +443,28 @@ class 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,
@@ -524,6 +632,21 @@ class cstruct:
         # signed long long: TypeAlias = int64
         # unsigned long long: TypeAlias = uint64
 
+        u8: TypeAlias = uint8
+        u16: TypeAlias = uint16
+        u32: TypeAlias = uint32
+        u64: TypeAlias = uint64
+        u128: TypeAlias = uint128
+        __u8: TypeAlias = uint8
+        __u16: TypeAlias = uint16
+        __u32: TypeAlias = uint32
+        __u64: TypeAlias = uint64
+        __u128: TypeAlias = uint128
+        uchar: TypeAlias = uint8
+        ushort: TypeAlias = uint16
+        uint: TypeAlias = uint32
+        ulong: TypeAlias = uint32
+
         BYTE: TypeAlias = uint8
         CHAR: TypeAlias = char
         SHORT: TypeAlias = int16
@@ -589,19 +712,17 @@ class cstruct:
         _QWORD: TypeAlias = uint64
         _OWORD: TypeAlias = uint128
 
-        u1: TypeAlias = uint8
-        u2: TypeAlias = uint16
-        u4: TypeAlias = uint32
-        u8: TypeAlias = uint64
-        u16: TypeAlias = uint128
-        __u8: TypeAlias = uint8
-        __u16: TypeAlias = uint16
-        __u32: TypeAlias = uint32
-        __u64: TypeAlias = uint64
-        uchar: TypeAlias = uint8
-        ushort: TypeAlias = uint16
-        uint: TypeAlias = uint32
-        ulong: TypeAlias = uint32
+
+# Attribute names that types and constants may never shadow: the public API and methods of the cstruct
+# class itself, plus the instance attributes created in __init__.
+_RESERVED_NAMES = frozenset(dir(cstruct)) | {
+    "endian",
+    "consts",
+    "types",
+    "includes",
+    "pointer",
+    "_anonymous_count",
+}
 
 
 def ctypes(structure: type[Structure]) -> type[_ctypes.Structure]:

dissect/cstruct/{exceptions.py → exception.py}

@@ -5,6 +5,10 @@ class Error(Exception):
     pass
 
 
+class LexerError(Error):
+    pass
+
+
 class ParserError(Error):
     pass
 
@@ -23,7 +27,3 @@ class ArraySizeError(Error):
 
 class ExpressionParserError(Error):
     pass
-
-
-class ExpressionTokenizerError(Error):
-    pass