dissect.cstruct 5.0.dev0__py3-none-any.whl → 5.0.dev7__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,13 +2,16 @@ 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
 from dissect.cstruct.types import (
@@ -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,25 +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.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),
@@ -102,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",
@@ -169,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
@@ -201,27 +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.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
 
@@ -230,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
 
@@ -264,8 +286,49 @@ 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.
 
@@ -301,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.
 
@@ -328,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.
@@ -351,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,
@@ -518,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
@@ -583,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/expression.py

@@ -2,9 +2,9 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from dissect.cstruct.exceptions import ExpressionParserError
+from dissect.cstruct.exception import ExpressionParserError
 from dissect.cstruct.lexer import IDENTIFIER_TYPES, Lexer, TokenCursor, TokenType
-from dissect.cstruct.utils import offsetof, sizeof
+from dissect.cstruct.util import offsetof, sizeof
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -12,6 +12,8 @@ if TYPE_CHECKING:
     from dissect.cstruct import cstruct
     from dissect.cstruct.lexer import Token
 
+_MISSING = object()
+
 BINARY_OPERATORS: dict[TokenType, Callable[[int, int], int]] = {
     TokenType.PIPE: lambda a, b: a | b,
     TokenType.CARET: lambda a, b: a ^ b,
@@ -164,16 +166,13 @@ class Expression(TokenCursor):
                             obj = context[part]
                         elif part in cs.consts:
                             obj = cs.consts[part]
-                        elif part in cs.typedefs:
-                            obj = cs.resolve(part)
+                        elif part in cs.types:
+                            obj = cs.types[part]
                         else:
                             raise self._error(f"Unknown identifier: '{ident}'", token=token)
                     else:
-                        if isinstance(obj, dict) and part in obj:
-                            obj = obj[part]
-                        elif hasattr(obj, part):
-                            obj = getattr(obj, part)
-                        else:
+                        obj = obj.get(part, _MISSING) if isinstance(obj, dict) else getattr(obj, part, _MISSING)
+                        if obj is _MISSING:
                             raise self._error(f"Unknown identifier: '{ident}'", token=token)
 
                 try:

dissect/cstruct/lexer.py

@@ -4,7 +4,7 @@ import enum
 import re
 from typing import TYPE_CHECKING
 
-from dissect.cstruct.exceptions import LexerError
+from dissect.cstruct.exception import LexerError
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -134,6 +134,14 @@ def tokenize(data: str) -> list[Token]:
     return Lexer(data).tokenize()
 
 
+def format_error_context(data: str, lineno: int, window: int = 1) -> str:
+    """Format a snippet of the input data around a given line number for error messages."""
+    start = max(0, lineno - window - 2)
+    context = data.splitlines()[start:lineno]
+    no_len = len(str(lineno)) + 2
+    return "\n".join(f"{start + i + 1:>{no_len}}: {line}" for i, line in enumerate(context))
+
+
 class Lexer:
     """Lexer compatible with C-like syntax for struct definitions and preprocessor directives."""
 
@@ -213,8 +221,10 @@ class Lexer:
 
         return self._take()
 
-    def _error(self, msg: str, *, line: int | None = None) -> LexerError:
-        return LexerError(f"line {line if line is not None else self._line}: {msg}")
+    def _error(self, msg: str, *, lineno: int | None = None) -> LexerError:
+        lineno = lineno if lineno is not None else self._line
+        content = format_error_context(self.data, lineno)
+        return LexerError(f"line {lineno}: {msg}\n{content}")
 
     def _emit(self, type: TokenType, value: str, line: int, column: int = 0) -> None:
         """Emit a token with the given type and value at the specified line and column."""
@@ -402,7 +412,7 @@ class Lexer:
         keyword = self._read_identifier()
 
         if (token_type := _PP_KEYWORDS.get(keyword)) is None:
-            raise self._error(f"unknown preprocessor directive '#{keyword}'", line=line)
+            raise self._error(f"unknown preprocessor directive '#{keyword}'", lineno=line)
 
         self._emit(token_type, keyword, line, col)
 
@@ -410,7 +420,7 @@ class Lexer:
             self._skip_whitespace_and_comments()
 
             if not (name := self._read_identifier()):
-                raise self._error("expected identifier after '#define'", line=line)
+                raise self._error("expected identifier after '#define'", lineno=line)
             self._emit(TokenType.IDENTIFIER, name, line)
 
             self._skip_whitespace_and_comments()
@@ -434,7 +444,7 @@ class Lexer:
                 self._expect(">")  # Consume closing `>`
                 value = f"<{value}>"
             else:
-                raise self._error("expected include path after '#include'", line=line)
+                raise self._error("expected include path after '#include'", lineno=line)
 
             self._emit(TokenType.STRING, value, line)
 
@@ -486,7 +496,7 @@ class Lexer:
                 self._emit(_SINGLE_CHARS[ch], self._take(), line, col)
 
             else:
-                raise self._error(f"unexpected character {ch!r}", line=line)
+                raise self._error(f"unexpected character {ch!r}", lineno=line)
 
         self._emit(TokenType.EOF, "", self._line, self._column)
         return self._tokens