dissect.cstruct 5.0.dev11__tar.gz → 5.0.dev13__tar.gz

{dissect_cstruct-5.0.dev11 → dissect_cstruct-5.0.dev13}/CHANGELOG.md

@@ -22,6 +22,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   - Faster paths for parsing primitive types from bytes-like objects, enum and flag lookups and flag operators.
   - Null-terminated strings are read in chunks from seekable streams.
 
+### Fixed
+
+- Structures and unions no longer share nested structures and arrays between their default values.
+- Unions can be created with default values and read if they have doubly nested structures or nested unions.
+- Writing a `bytearray` to a `char` array no longer modifies the `bytearray`.
+- Writing a structure correctly returns the amount of bytes written, including padding and bit fields.
+- Writing a structure raises an `AttributeError` for missing attributes and for values of a nested structure field that aren't a structure, instead of writing default values.
+- A surrogate pair takes up two entries of a `wchar` array, so writing it no longer writes more bytes than the size of the array.
+- A `wchar` string can hold a lone surrogate, instead of raising an error for the whole structure when reading or writing it. Bytes are written back as they were stored.
+- Writing more than a single character to a `char` or `wchar` raises a `ValueError`, instead of writing all of it and moving the fields after it. An empty value is padded.
+
 ## [4.5] - 20-05-2025
 
 ### Added

{dissect_cstruct-5.0.dev11 → dissect_cstruct-5.0.dev13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dissect.cstruct
-Version: 5.0.dev11
+Version: 5.0.dev13
 Summary: A Dissect module implementing a parser for C-like structures: structure parsing in Python made easy
 Author-email: Dissect Team <dissect@fox-it.com>
 License-Expression: Apache-2.0

{dissect_cstruct-5.0.dev11 → dissect_cstruct-5.0.dev13}/dissect/cstruct/compiler.py

@@ -182,8 +182,10 @@ _GLOBALS = {
 }
 
 
-class _ReadSourceGenerator:
-    """Generates the source code of an optimized ``_read`` method for a structure or union."""
+class _SourceGenerator:
+    """Base class for generating the source code of an optimized method of a structure or union."""
+
+    name: str
 
     def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
         self.cls = cls
@@ -191,24 +193,17 @@ class _ReadSourceGenerator:
         self.align = cls.__align__
         self.inline = inline and not self.align
 
+        # The values available to the generated code
+        self.symbols: dict[str, Any] = {"cs": cls.__cs__}
+        self._symbol_names: dict[int, str] = {}
+
         # Witnesses of inlined structure types (symbol of the type -> symbol of its __init__ at compile time)
         self._witnesses: dict[str, str] = {}
-        self._fallback_symbol: str | None = None
 
         # The state of the bit buffer: the type of its current value and the remaining bits
         self._bits: tuple[type[BaseType], int] | None = None
         self._var_ids: dict[str, int] = {}
 
-        # The values available to the generated code
-        self.symbols: dict[str, Any] = {"cs": cls.__cs__}
-        self._symbol_names: dict[int, str] = {}
-
-        # (field name, local variable name) of all fields assigned so far
-        self._assigned: list[tuple[str, str]] = []
-        self._has_context = False
-        self._uses_start = False
-        self._uses_sizes = False
-
     def _type_symbol(self, type_: type[BaseType]) -> str:
         """Add a type as symbol, returns its name."""
         return self._add_symbol(type_, "type")
@@ -220,39 +215,12 @@ class _ReadSourceGenerator:
             self.symbols[symbol] = value
         return symbol
 
-    def _context(self, lines: list[str]) -> str:
-        """Return the name of the context dictionary, creating it from the fields read so far if needed."""
-        if not self._has_context:
-            self._has_context = True
-            items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
-            lines.append(f"values = {{{items}}}")
-
-        return "values"
-
     def _var(self, name: str) -> str:
         """Return a unique variable name for the generated code."""
         index = self._var_ids.get(name, 0)
         self._var_ids[name] = index + 1
         return f"_{name}{index}"
 
-    def _assign(self, field: Field, value: str) -> str:
-        """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
-        var = f"_value{len(self._assigned)}"
-        self._assigned.append((field._name, var))
-        if self._has_context:
-            return f"{var} = values[{field._name!r}] = {value}"
-        return f"{var} = {value}"
-
-    def _assign_local(self, field: Field, value: str) -> list[str]:
-        return [self._assign(field, value)]
-
-    def _fallback(self) -> str:
-        """Return the symbol of the non-inlined fallback read function."""
-        if self._fallback_symbol is None:
-            fallback = _Fallback(self.cls, self.fields)
-            self._fallback_symbol = self._add_symbol(fallback, "fallback")
-        return self._fallback_symbol
-
     def _witness(self, type_: type[Structure]) -> str:
         """Add an inlined type as symbol, with a witness to detect modifications to it after compilation."""
         symbol = self._type_symbol(type_)
@@ -263,10 +231,22 @@ class _ReadSourceGenerator:
         """Return the conditions that are true if any of the inlined types were modified after compilation."""
         return [f"{type_}.__init__ is not {init}" for type_, init in self._witnesses.items()]
 
+    def _starts_bit_value(self, field: Field) -> bool:
+        """Return whether a bit field is the first of a new value of the bit buffer."""
+        return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bit_buffer_type(field)
+
+    def _advance_bits(self, field: Field) -> int:
+        """Update the state of the bit buffer for a bit field, returns the bits that were remaining before it."""
+        if self._starts_bit_value(field):
+            self._bits = (_get_bit_buffer_type(field), _get_bits_type(field).__size__ * 8)
+        remaining = self._bits[1]
+        self._bits = (self._bits[0], remaining - field.bits)
+        return remaining
+
     def generate(self) -> MethodType:
         source = self.generate_source()
 
-        code = python_compile(source, f"<compiled {self.cls.__name__}._read>", "exec")
+        code = python_compile(source, f"<compiled {self.cls.__name__}.{self.name}>", "exec")
         exec(code, {**_GLOBALS, **self.symbols}, d := {})
         obj = d.popitem()[1]
         obj.__source__ = source
@@ -278,6 +258,56 @@ class _ReadSourceGenerator:
             return self._generate_union_source()
         return self._generate_struct_source()
 
+    def _generate_struct_source(self) -> str:
+        raise NotImplementedError
+
+    def _generate_union_source(self) -> str:
+        raise NotImplementedError
+
+
+class _ReadSourceGenerator(_SourceGenerator):
+    """Generates the source code of an optimized ``_read`` method for a structure or union."""
+
+    name = "_read"
+
+    def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
+        super().__init__(cls, fields, inline)
+
+        self._fallback_symbol: str | None = None
+
+        # (field name, local variable name) of all fields assigned so far
+        self._assigned: list[tuple[str, str]] = []
+        self._has_context = False
+        self._uses_start = False
+        self._uses_sizes = False
+
+    def _context(self, lines: list[str]) -> str:
+        """Return the name of the context dictionary, creating it from the fields read so far if needed."""
+        if not self._has_context:
+            self._has_context = True
+            items = ", ".join(f"{name!r}: {var}" for name, var in self._assigned)
+            lines.append(f"values = {{{items}}}")
+
+        return "values"
+
+    def _assign(self, field: Field, value: str) -> str:
+        """Assign a value to a new local variable for the given field, and the context dictionary if it exists."""
+        var = f"_value{len(self._assigned)}"
+        self._assigned.append((field._name, var))
+        if self._has_context:
+            return f"{var} = values[{field._name!r}] = {value}"
+        return f"{var} = {value}"
+
+    def _assign_local(self, field: Field, value: str) -> list[str]:
+        return [self._assign(field, value)]
+
+    def _fallback(self) -> str:
+        """Return the symbol of the non-inlined fallback read function."""
+        if self._fallback_symbol is None:
+            fallback = _Fallback(self.cls, self.fields)
+            self._fallback_symbol = self._add_symbol(fallback, "fallback")
+        return self._fallback_symbol
+
     def _generate_struct_source(self) -> str:
         """Generate the code for reading a structure."""
         read_code = "\n".join(self._generate_fields_source())
@@ -405,6 +435,16 @@ class _ReadSourceGenerator:
 
             # Everything else - basic and composite types (and arrays of them), and inlined structures
             else:
+                if field.offset is None:
+                    # The alignment depends on the position in the stream, so this needs a new block
+                    if self.align:
+                        yield from flush()
+                # A block is read in order, so a field before the current position needs a new block
+                elif current_offset is None or field.offset < current_offset:
+                    yield from flush()
+                elif current_block:
+                    current_offset = field.offset
+
                 if not current_block:
                     yield from align_to_field(field)
                 current_block.append(field)
@@ -456,18 +496,11 @@ class _ReadSourceGenerator:
         element = self._type_symbol(field.type.__type__)
         return f"_type_call({type_}, {element}._read_array(stream, _count, context={context}, endian=endian))"
 
-    def _starts_bit_value(self, field: Field) -> bool:
-        """Return whether a bit field is read from a new value of the bit buffer."""
-        return self._bits is None or self._bits[1] == 0 or self._bits[0] is not _get_bits_type(field)
-
     def _generate_bits_source(self, field: Field) -> Iterator[str]:
         """Generate the code for reading a bit field."""
         read_type = _get_bits_type(field)
         new_value = self._starts_bit_value(field)
-        if new_value:
-            self._bits = (read_type, read_type.__size__ * 8)
-        remaining = self._bits[1]
-        self._bits = (read_type, remaining - field.bits)
+        remaining = self._advance_bits(field)
 
         lines = []
         if new_value:
@@ -498,7 +531,7 @@ class _ReadSourceGenerator:
         """Generate the code for reading consecutive packed fields."""
         unpack = _Unpack("data")
         reads = []
-        size = self._pack_fields(fields, 0, unpack, reads, self._assign_local, self.align)
+        size = self._pack_fields(fields, 0, unpack, reads, self._assign_local)
 
         lines = [f"buf = stream.read({size})", _check_eof(size)]
         if unpack.has_values():
@@ -513,7 +546,6 @@ class _ReadSourceGenerator:
         unpack: _Unpack,
         lines: list[str],
         target: Callable[[Field, str], list[str]],
-        align: bool = False,
     ) -> int:
         """Generate the code for reading the given fields from ``buf``, starting at the given offset.
 
@@ -523,7 +555,7 @@ class _ReadSourceGenerator:
         Returns the offset after the last field.
         """
         inline = _inline_kind if self.inline else None
-        for field, count, char in _generate_struct_info(fields, align, inline=inline):
+        for field, count, char in _generate_struct_info(fields, inline=inline):
             if field is None:
                 # Padding
                 unpack.add(count, "x")
@@ -595,16 +627,15 @@ class _ReadSourceGenerator:
         lines.append(f"_object_setattr({var}, '_endian', endian)")
         self._proxify(var, var, fields, lines)
 
-    def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str]) -> None:
+    def _proxify(self, union: str, var: str, fields: list[Field], lines: list[str], attr: str | None = None) -> None:
         """Generate the code equivalent to ``Union._proxify``."""
         for field in fields:
             if issubclass(field.type, Structure):
                 nested = self._var("nested")
+                member = attr or field._name
                 lines.append(f"{nested} = _getattr({var}, {field._name!r})")
-                lines.append(
-                    f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {field._name!r}, {nested}))"
-                )
-                self._proxify(union, nested, field.type.__fields__, lines)
+                lines.append(f"_object_setattr({var}, {field._name!r}, _UnionProxy({union}, {member!r}, {nested}))")
+                self._proxify(union, nested, field.type.__fields__, lines, member)
 
     def _generate_packed_field_source(
         self,
@@ -622,19 +653,18 @@ class _ReadSourceGenerator:
             array = self._type_symbol(field.type)
             element = self._type_symbol(field_type.__type__)
 
-            if issubclass(field_type.__type__, Int):
-                elem_size = field_type.__type__.__size__
-                signed = bool(field_type.__type__.signed)
+            if issubclass(read_type, Int):
+                # Everything that's read as an integer type is sliced directly out of the buffer
+                elem_size = read_type.__size__
+                signed = bool(read_type.signed)
                 lines.append(f"_array_buf = {getter}")
                 raw = f"_int_from_bytes(_array_buf[i:i + {elem_size}], _byteorder[endian], signed={signed})"
-                value = f"[_type_call({element}, {raw}) for i in range(0, {count}, {elem_size})]"
-            elif issubclass(field_type.__type__, Pointer):
+                getter = f"[{raw} for i in range(0, {count}, {elem_size})]"
+
+            if issubclass(field_type.__type__, Pointer):
                 context = self._context(lines)
                 item = f"{element}.__new__({element}, value, stream, context={context}, endian=endian)"
                 value = f"[{item} for value in {getter}]"
-            elif issubclass(read_type, (Wchar, Int)):
-                # Preserve the behaviour of the generic parsing path for these types
-                value = f"[{element}(value) for value in {getter}]"
             else:
                 value = f"list(map(_type_call, _repeat({element}), {getter}))"
 
@@ -647,7 +677,7 @@ class _ReadSourceGenerator:
             encoding = self._add_symbol(read_type.__encoding_map__, "encoding")
             # Using str() instead of .decode() also supports other bytes-like objects returned by the stream
             type_ = self._type_symbol(field.type)
-            lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian]))"))
+            lines.extend(target(field, f"_type_call({type_}, str({getter}, {encoding}[endian], 'surrogatepass'))"))
 
         elif issubclass(field_type, Pointer):
             context = self._context(lines)
@@ -694,8 +724,7 @@ class _Unpack:
 
     def functions(self, name: str) -> dict[str, Callable]:
         """Return the given function of the struct for every endianness."""
-        fmt = _optimize_struct_fmt(self.items)
-        return {endian: getattr(_struct(endian, fmt), name) for endian in _ENDIANNESS}
+        return _struct_functions(_optimize_struct_fmt(self.items), name)
 
 
 def _check_union(cls: type[Union], fields: list[Field], inline: bool = True) -> int:
@@ -809,25 +838,18 @@ def _block(lines: list[str]) -> str:
 
 def _generate_struct_info(
     fields: list[Field],
-    align: bool = False,
     inline: Callable[[type[BaseType]], str | None] | None = None,
 ) -> Iterator[tuple[Field, int, str]]:
     if not fields:
         return
 
     current_offset = fields[0].offset
-    imaginary_offset = 0
     for field in fields:
         # We moved -- probably due to alignment
         if field.offset is not None and (drift := field.offset - current_offset) > 0:
             yield None, drift, "x"
             current_offset += drift
 
-        if align and field.offset is None and (drift := -imaginary_offset & (field.alignment - 1)) > 0:
-            # Assume we started at a correctly aligned boundary
-            yield None, drift, "x"
-            imaginary_offset += drift
-
         count = 1
         read_type = _get_read_type(field.type)
 
@@ -853,25 +875,41 @@ def _generate_struct_info(
         elif issubclass(read_type, (Char, Wchar, Int)):
             yield field, count * read_type.__size__, "x"
 
-        size = count * read_type.__size__
-        imaginary_offset += size
         if current_offset is not None:
-            current_offset += size
+            current_offset += count * read_type.__size__
 
 
 def _optimize_struct_fmt(items: Iterable[tuple[int, str]]) -> str:
-    """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char."""
+    """Create a struct format string from ``(count, char)`` items, merging consecutive items of the same char.
+
+    Byte strings aren't merged, since those are a single value regardless of their count.
+    """
     chars = []
     for char, group in groupby(items, key=itemgetter(1)):
-        if count := sum(count for count, _ in group):
+        if char == "s":
+            chars.extend(f"{count}{char}" for count, _ in group)
+        elif count := sum(count for count, _ in group):
             chars.append(f"{count if count > 1 else ''}{char}")
 
     return "".join(chars)
 
 
+def _struct_functions(fmt: str, name: str) -> dict[str, Callable]:
+    """Return the given function of a struct with the given format for every endianness.
+
+    The ``@`` endianness uses the ``=`` format, since ``@`` would add padding between items to align them.
+    """
+    return {endian: getattr(_struct("=" if endian == "@" else endian, fmt), name) for endian in _ENDIANNESS}
+
+
+def _get_bit_buffer_type(field: Field) -> type[BaseType]:
+    """Return the type that ``BitBuffer`` uses for a bit field, only bit fields of the same type share a value."""
+    return field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
+
+
 def _get_bits_type(field: Field) -> type[BaseType]:
     """Return the type that ``BitBuffer`` reads a bit field as."""
-    type_ = field.type.__type__ if issubclass(field.type, (Enum, Flag)) else field.type
+    type_ = _get_bit_buffer_type(field)
     return type_.__cs__.uint8 if issubclass(type_, Char) else type_
 
 

{dissect_cstruct-5.0.dev11 → dissect_cstruct-5.0.dev13}/dissect/cstruct/cstruct.py

@@ -86,7 +86,7 @@ class cstruct:
             "int24": self._make_int_type("int24", 3, True, alignment=4),
             "uint24": self._make_int_type("uint24", 3, False, alignment=4),
             "int48": self._make_int_type("int48", 6, True, alignment=8),
-            "uint48": self._make_int_type("int48", 6, False, alignment=8),
+            "uint48": self._make_int_type("uint48", 6, False, alignment=8),
             "int128": self._make_int_type("int128", 16, True, alignment=16),
             "uint128": self._make_int_type("uint128", 16, False, alignment=16),
 

{dissect_cstruct-5.0.dev11 → dissect_cstruct-5.0.dev13}/dissect/cstruct/types/base.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import functools
 import warnings
 from io import SEEK_CUR, BytesIO
+from types import MethodType
 from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar, TypeVar
 
 from dissect.cstruct.exception import ArraySizeError
@@ -19,6 +20,8 @@ if TYPE_CHECKING:
 
 EOF = -0xE0F  # Negative counts are illegal anyway, so abuse that for our EOF sentinel
 
+_IMMUTABLE_TYPES = (int, float, bytes, str)
+
 
 class MetaType(type):
     """Base metaclass for cstruct type classes."""
@@ -218,7 +221,8 @@ class MetaType(type):
             array: The array to write.
             endian: The endianness to use when parsing (little, big, network, <, >, !, @ or =).
         """
-        return sum(cls._write(stream, entry, endian=endian) for entry in array)
+        write = cls._write
+        return sum(write(stream, entry, endian=endian) for entry in array)
 
     def _write_0(cls, stream: BinaryIO, array: list[Self], *, endian: Endianness) -> int:  # type: ignore
         """Internal function for writing null-terminated arrays.
@@ -295,20 +299,29 @@ class _overload:
         b'\\x7b\\x00\\x00\\x00'
     """
 
-    def __init__(self, func: Callable[..., Any]) -> None:
+    def __init__(self, func: Callable[..., Any], method: Callable[..., Any]) -> None:
         self.func = func
+        self.method = method
 
-    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[[], bytes]:
+    def __get__(self, instance: BaseType | None, owner: type[BaseType]) -> Callable[..., Any]:
         if instance is None:
             return functools.partial(self.func, owner)
-        return functools.partial(self.func, instance.__class__, value=instance)
+        return MethodType(self.method, instance)
+
+
+def _instance_dumps(self: BaseType, *, endian: AllowedEndianness | None = None) -> bytes:
+    return MetaType.dumps(self.__class__, self, endian=endian)
+
+
+def _instance_write(self: BaseType, stream: BinaryIO, *, endian: AllowedEndianness | None = None) -> int:
+    return MetaType.write(self.__class__, stream, self, endian=endian)
 
 
 class BaseType(metaclass=MetaType):
     """Base class for cstruct type classes."""
 
-    dumps = _overload(MetaType.dumps)
-    write = _overload(MetaType.write)
+    dumps = _overload(MetaType.dumps, _instance_dumps)
+    write = _overload(MetaType.write, _instance_write)
 
     def __len__(self) -> int:
         """Return the byte size of the type."""
@@ -338,9 +351,10 @@ class BaseArray(BaseType):
 
     @classmethod
     def __default__(cls) -> BaseType:
-        return type.__call__(
-            cls, [cls.__type__.__default__()] * (cls.__num_entries__ if isinstance(cls.__num_entries__, int) else 0)
-        )
+        count = cls.__num_entries__ if isinstance(cls.__num_entries__, int) else 0
+        if isinstance(value := cls.__type__.__default__(), _IMMUTABLE_TYPES):
+            return type.__call__(cls, [value] * count)
+        return type.__call__(cls, [cls.__type__.__default__() for _ in range(count)])
 
     @classmethod
     def _read(cls, stream: BinaryIO, *, context: dict[str, Any] | None = None, endian: Endianness) -> list[BaseType]:

{dissect_cstruct-5.0.dev11 → dissect_cstruct-5.0.dev13}/dissect/cstruct/types/char.py

@@ -23,8 +23,9 @@ class CharArray(bytes, BaseArray):
         return type.__call__(cls, super()._read(stream, context=context, endian=endian))
 
     @classmethod
-    def _write(cls, stream: BinaryIO, data: bytes | str | list[int], *, endian: Endianness) -> int:
-        if isinstance(data, list):
+    def _write(cls, stream: BinaryIO, data: bytes | bytearray | str | list[int], *, endian: Endianness) -> int:
+        # Copy mutable data types to bytes so we don't modify them in place
+        if isinstance(data, (list, bytearray)):
             data = bytes(data)
         elif isinstance(data, str):
             data = data.encode("latin-1")
@@ -92,4 +93,7 @@ class Char(bytes, BaseType):
         if isinstance(data, str):
             data = data.encode("latin-1")
 
-        return stream.write(data)
+        if (actual_size := len(data)) > 1:
+            raise ValueError(f"Expected a single character, got {actual_size} instead.")
+
+        return stream.write(data or b"\x00")