dissect.cstruct 5.0.dev13__tar.gz → 5.0.dev14__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/CHANGELOG.md +3 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/PKG-INFO +1 -1
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/compiler.py +440 -15
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/base.py +14 -2
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/PKG-INFO +1 -1
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/scm_version.json +2 -2
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_align.py +52 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_bitfield.py +36 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_compiler.py +605 -3
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_base.py +55 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_char.py +72 -2
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_int.py +26 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_packed.py +25 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_pointer.py +21 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_structure.py +222 -1
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_union.py +119 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_void.py +15 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_wchar.py +22 -1
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/.git-blame-ignore-revs +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/.gitattributes +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/COPYRIGHT +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/LICENSE +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/MANIFEST.in +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/README.md +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/__init__.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/bitbuffer.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/cstruct.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/exception.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/expression.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/lexer.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/parser.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/tools/__init__.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/tools/stubgen.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/__init__.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/char.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/enum.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/flag.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/int.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/leb128.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/packed.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/pointer.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/structure.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/void.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/types/wchar.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect/cstruct/util.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/SOURCES.txt +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/dependency_links.txt +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/entry_points.txt +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/requires.txt +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/scm_file_list.json +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/dissect.cstruct.egg-info/top_level.txt +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/examples/disk.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/examples/mirai.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/examples/pe.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/examples/protobuf.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/examples/secdesc.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/pyproject.toml +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/setup.cfg +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/__init__.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/_data/testdef.txt +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/_docs/Makefile +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/_docs/__init__.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/_docs/conf.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/_docs/index.rst +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/conftest.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_annotations.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_basic.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_benchmark.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_bitbuffer.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_ctypes.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_expression.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_lexer.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_parser.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_tools_stubgen.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_custom.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_enum.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_flag.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_types_leb128.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/test_utils.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tests/utils.py +0 -0
- {dissect_cstruct-5.0.dev13 → dissect_cstruct-5.0.dev14}/tox.ini +0 -0
|
@@ -21,6 +21,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
|
|
|
21
21
|
- Compiled structures calculate the sizes of dynamic arrays directly, instead of evaluating the expression.
|
|
22
22
|
- Faster paths for parsing primitive types from bytes-like objects, enum and flag lookups and flag operators.
|
|
23
23
|
- Null-terminated strings are read in chunks from seekable streams.
|
|
24
|
+
- Improve writing performance:
|
|
25
|
+
- The compiler now generates code for writing structures and unions too! It uses a similar approach, grouping fields and writes where possible.
|
|
26
|
+
- `obj.dumps()` and `obj.write()` no longer create a `functools.partial` on every access. As a result, they no longer accept a `value` argument.
|
|
24
27
|
|
|
25
28
|
### Fixed
|
|
26
29
|
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: dissect.cstruct
|
|
3
|
-
Version: 5.0.
|
|
3
|
+
Version: 5.0.dev14
|
|
4
4
|
Summary: A Dissect module implementing a parser for C-like structures: structure parsing in Python made easy
|
|
5
5
|
Author-email: Dissect Team <dissect@fox-it.com>
|
|
6
6
|
License-Expression: Apache-2.0
|
|
@@ -4,6 +4,7 @@ from __future__ import annotations
|
|
|
4
4
|
|
|
5
5
|
import io
|
|
6
6
|
import logging
|
|
7
|
+
import struct
|
|
7
8
|
from enum import Enum
|
|
8
9
|
from itertools import groupby, repeat
|
|
9
10
|
from keyword import iskeyword
|
|
@@ -11,7 +12,7 @@ from operator import itemgetter
|
|
|
11
12
|
from textwrap import indent
|
|
12
13
|
from typing import TYPE_CHECKING, Any
|
|
13
14
|
|
|
14
|
-
from dissect.cstruct.exception import ExpressionParserError
|
|
15
|
+
from dissect.cstruct.exception import ArraySizeError, ExpressionParserError
|
|
15
16
|
from dissect.cstruct.expression import SOURCE_GLOBALS, Expression
|
|
16
17
|
from dissect.cstruct.types import (
|
|
17
18
|
Array,
|
|
@@ -30,7 +31,7 @@ from dissect.cstruct.types import (
|
|
|
30
31
|
WcharArray,
|
|
31
32
|
)
|
|
32
33
|
from dissect.cstruct.types.enum import EnumMetaType
|
|
33
|
-
from dissect.cstruct.types.structure import UnionMetaType, UnionProxy
|
|
34
|
+
from dissect.cstruct.types.structure import StructureMetaType, UnionMetaType, UnionProxy
|
|
34
35
|
from dissect.cstruct.util import ENDIANNESS_TO_BYTEORDER_MAP, _struct
|
|
35
36
|
|
|
36
37
|
if TYPE_CHECKING:
|
|
@@ -63,25 +64,33 @@ python_compile = compile
|
|
|
63
64
|
|
|
64
65
|
|
|
65
66
|
def compile_structure(cls: type[Structure]) -> type[Structure]:
|
|
66
|
-
"""Compile the ``_read``
|
|
67
|
+
"""Compile the ``_read`` and ``_write`` methods of a structure or union type, if the compiler supports it.
|
|
67
68
|
|
|
68
69
|
The compilation happens on first use.
|
|
69
70
|
"""
|
|
70
71
|
try:
|
|
71
72
|
_check_supported(cls, cls.__fields__)
|
|
72
73
|
cls._read = _LazyCompileRead(cls)
|
|
74
|
+
cls._write = _LazyCompileWrite(cls)
|
|
73
75
|
cls.__compiled__ = True
|
|
74
76
|
except Exception as e:
|
|
75
|
-
# Revert to the interpreted reader
|
|
77
|
+
# Revert to the interpreted reader and writer
|
|
76
78
|
log.debug("Failed to compile %s", cls, exc_info=e)
|
|
77
|
-
|
|
78
|
-
cls.__compiled__ = False
|
|
79
|
+
_revert_to_interpreted(cls)
|
|
79
80
|
|
|
80
81
|
return cls
|
|
81
82
|
|
|
82
83
|
|
|
83
|
-
|
|
84
|
-
|
|
84
|
+
def _revert_to_interpreted(cls: type[Structure]) -> None:
|
|
85
|
+
cls._read = classmethod(type(cls)._read)
|
|
86
|
+
cls._write = classmethod(type(cls)._write)
|
|
87
|
+
cls.__compiled__ = False
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
class _LazyCompile:
|
|
91
|
+
"""Descriptor that compiles a method of a structure on first access and replaces itself with it."""
|
|
92
|
+
|
|
93
|
+
name: str
|
|
85
94
|
|
|
86
95
|
def __init__(self, cls: type[Structure]):
|
|
87
96
|
self.cls = cls
|
|
@@ -103,18 +112,47 @@ class _LazyCompileRead:
|
|
|
103
112
|
inline = all(type_.__dict__.get("__init__") is init for type_, init in self.snapshot.items())
|
|
104
113
|
|
|
105
114
|
try:
|
|
106
|
-
self.method =
|
|
115
|
+
self.method = self._generator(inline).generate()
|
|
107
116
|
except Exception as e:
|
|
108
|
-
# Revert to the slower
|
|
117
|
+
# Revert to the slower interpreted method
|
|
109
118
|
log.debug("Failed to compile %s", cls, exc_info=e)
|
|
110
|
-
self.method =
|
|
111
|
-
cls.__compiled__ = False
|
|
119
|
+
self.method = self._revert()
|
|
112
120
|
|
|
113
|
-
if cls.__dict__.get(
|
|
114
|
-
cls.
|
|
121
|
+
if cls.__dict__.get(self.name) is self:
|
|
122
|
+
setattr(cls, self.name, self.method)
|
|
115
123
|
|
|
116
124
|
return self.method
|
|
117
125
|
|
|
126
|
+
def _generator(self, inline: bool) -> _SourceGenerator:
|
|
127
|
+
raise NotImplementedError
|
|
128
|
+
|
|
129
|
+
def _revert(self) -> classmethod:
|
|
130
|
+
"""Return the interpreted version of the method."""
|
|
131
|
+
return classmethod(getattr(type(self.cls), self.name))
|
|
132
|
+
|
|
133
|
+
|
|
134
|
+
class _LazyCompileRead(_LazyCompile):
|
|
135
|
+
"""Descriptor that compiles the ``_read`` method of a structure on first access and replaces itself with it."""
|
|
136
|
+
|
|
137
|
+
name = "_read"
|
|
138
|
+
|
|
139
|
+
def _generator(self, inline: bool) -> _SourceGenerator:
|
|
140
|
+
return _ReadSourceGenerator(self.cls, self.fields, inline)
|
|
141
|
+
|
|
142
|
+
def _revert(self) -> classmethod:
|
|
143
|
+
# The structure no longer counts as compiled, so adding a field won't recompile _write either
|
|
144
|
+
_revert_to_interpreted(self.cls)
|
|
145
|
+
return super()._revert()
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
class _LazyCompileWrite(_LazyCompile):
|
|
149
|
+
"""Descriptor that compiles the ``_write`` method of a structure on first access and replaces itself with it."""
|
|
150
|
+
|
|
151
|
+
name = "_write"
|
|
152
|
+
|
|
153
|
+
def _generator(self, inline: bool) -> _SourceGenerator:
|
|
154
|
+
return _WriteSourceGenerator(self.cls, self.fields, inline)
|
|
155
|
+
|
|
118
156
|
|
|
119
157
|
def _snapshot_structures(fields: list[Field], result: dict[type, Any]) -> dict[type, Any]:
|
|
120
158
|
for field in fields:
|
|
@@ -166,7 +204,29 @@ _ENDIANNESS = ("<", ">", "!", "@", "=")
|
|
|
166
204
|
# Types of which the value can be used as an integer directly
|
|
167
205
|
_INTEGER_TYPES = (Int, Packed, Enum, Flag)
|
|
168
206
|
|
|
169
|
-
#
|
|
207
|
+
# Types of which the value can be packed, unless they have a custom way of writing
|
|
208
|
+
_PACKABLE_TYPES = (Array, Char, CharArray, Int, Packed, Pointer, Void, VoidArray, Wchar, WcharArray)
|
|
209
|
+
|
|
210
|
+
|
|
211
|
+
def _field_bytes(type_: type[BaseType], value: Any, endian: str) -> bytes:
|
|
212
|
+
"""Return the bytes of a char or wchar field that isn't of the right type yet or doesn't fit."""
|
|
213
|
+
buf = type_.dumps(type_.__default__() if value is None else value, endian=endian)
|
|
214
|
+
if len(buf) > type_.__size__:
|
|
215
|
+
raise ArraySizeError(f"Expected {type_.__size__} bytes for {type_.__name__}, got {len(buf)} instead.")
|
|
216
|
+
return buf
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
def _array_values(type_: type[Array], value: list[Any] | None) -> list[Any]:
|
|
220
|
+
"""Return the values of an array field that doesn't have the right size yet."""
|
|
221
|
+
return type_.__default__() if value is None else type_._pad(value)
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def _or_default(value: Any, type_: type[BaseType]) -> Any:
|
|
225
|
+
"""Return the value of a field, or the default of its type if it's ``None``."""
|
|
226
|
+
return type_.__default__() if value is None else value
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
# Globals that are available to all generated functions
|
|
170
230
|
_GLOBALS = {
|
|
171
231
|
"_type_call": type.__call__,
|
|
172
232
|
"_object_new": object.__new__,
|
|
@@ -178,6 +238,12 @@ _GLOBALS = {
|
|
|
178
238
|
"_getattr": getattr,
|
|
179
239
|
"_UnionProxy": UnionProxy,
|
|
180
240
|
"_union_read": UnionMetaType._read,
|
|
241
|
+
"_struct_write": StructureMetaType._write,
|
|
242
|
+
"_union_write": UnionMetaType._write,
|
|
243
|
+
"_struct_error": struct.error,
|
|
244
|
+
"_field_bytes": _field_bytes,
|
|
245
|
+
"_array_values": _array_values,
|
|
246
|
+
"_or_default": _or_default,
|
|
181
247
|
**SOURCE_GLOBALS,
|
|
182
248
|
}
|
|
183
249
|
|
|
@@ -727,6 +793,319 @@ class _Unpack:
|
|
|
727
793
|
return _struct_functions(_optimize_struct_fmt(self.items), name)
|
|
728
794
|
|
|
729
795
|
|
|
796
|
+
class _WriteSourceGenerator(_SourceGenerator):
|
|
797
|
+
"""Generates the source code of an optimized ``_write`` method for a structure or union."""
|
|
798
|
+
|
|
799
|
+
name = "_write"
|
|
800
|
+
|
|
801
|
+
def __init__(self, cls: type[Structure], fields: list[Field], inline: bool = True):
|
|
802
|
+
super().__init__(cls, fields, inline)
|
|
803
|
+
|
|
804
|
+
self._uses_start = False
|
|
805
|
+
|
|
806
|
+
def _generate_method_source(self, interpreted: str, lines: Iterable[str]) -> str:
|
|
807
|
+
"""Generate the code of the method, ``interpreted`` is the name of the method to fall back to."""
|
|
808
|
+
write_code = "\n".join(lines)
|
|
809
|
+
|
|
810
|
+
# Fall back to the generic method if any of the inlined types was modified after compilation
|
|
811
|
+
checks = [f"cls is not {self._add_symbol(self.cls, 'cls')}", *self._modified_checks()]
|
|
812
|
+
|
|
813
|
+
preamble = [f"if {' or '.join(checks)}: return {interpreted}(cls, stream, data, endian=endian)"]
|
|
814
|
+
if self._uses_start:
|
|
815
|
+
preamble.append("start = stream.tell()")
|
|
816
|
+
|
|
817
|
+
code = indent("\n".join([*preamble, write_code]), " ")
|
|
818
|
+
|
|
819
|
+
return f"def _write(cls, stream, data, *, endian):\n{code}"
|
|
820
|
+
|
|
821
|
+
def _generate_struct_source(self) -> str:
|
|
822
|
+
"""Generate the code for writing a structure."""
|
|
823
|
+
return self._generate_method_source("_struct_write", list(self._generate_fields_source()))
|
|
824
|
+
|
|
825
|
+
def _generate_union_source(self) -> str:
|
|
826
|
+
"""Generate the code for writing a union."""
|
|
827
|
+
if not _can_pack(self.cls):
|
|
828
|
+
raise TypeError(f"Unsupported union for compiler: {self.cls}")
|
|
829
|
+
|
|
830
|
+
pack = _Pack()
|
|
831
|
+
self._pack_union("data", self.cls, pack)
|
|
832
|
+
|
|
833
|
+
lines = [self._generate_pack_source(pack), f"return {self.cls.__size__}"]
|
|
834
|
+
return self._generate_method_source("_union_write", lines)
|
|
835
|
+
|
|
836
|
+
def _generate_fields_source(self) -> Iterator[str]:
|
|
837
|
+
"""Generate the code for writing all fields."""
|
|
838
|
+
# The offset of the current field in the structure, None if it's unknown due to a dynamic field
|
|
839
|
+
current_offset = 0
|
|
840
|
+
current_block: list[Field] = []
|
|
841
|
+
current_bits: list[Field] = []
|
|
842
|
+
|
|
843
|
+
def flush() -> Iterator[str]:
|
|
844
|
+
if current_block:
|
|
845
|
+
yield from self._generate_packed_source(current_block)
|
|
846
|
+
current_block[:] = []
|
|
847
|
+
if current_bits:
|
|
848
|
+
yield from self._generate_bits_source(current_bits)
|
|
849
|
+
current_bits[:] = []
|
|
850
|
+
|
|
851
|
+
def align_to_field(field: Field) -> Iterator[str]:
|
|
852
|
+
nonlocal current_offset
|
|
853
|
+
|
|
854
|
+
if field.offset is not None and field.offset != current_offset:
|
|
855
|
+
if current_offset is None:
|
|
856
|
+
# This only writes padding if the stream is before the offset, otherwise the amount is negative
|
|
857
|
+
yield f'stream.write(b"\\x00" * (start + {field.offset} - stream.tell()))'
|
|
858
|
+
else:
|
|
859
|
+
yield f"stream.write({bytes(field.offset - current_offset)!r})"
|
|
860
|
+
current_offset = field.offset
|
|
861
|
+
|
|
862
|
+
if self.align and field.offset is None:
|
|
863
|
+
yield f'stream.write(b"\\x00" * (-stream.tell() & ({field.alignment} - 1)))'
|
|
864
|
+
|
|
865
|
+
for field in self.fields:
|
|
866
|
+
if not field.bits:
|
|
867
|
+
self._bits = None
|
|
868
|
+
if current_bits:
|
|
869
|
+
yield from flush()
|
|
870
|
+
|
|
871
|
+
# Unlike reading, we can't move back in the stream
|
|
872
|
+
if None not in (field.offset, current_offset) and field.offset < current_offset:
|
|
873
|
+
raise TypeError(f"Unsupported offset for compiler: {field}")
|
|
874
|
+
|
|
875
|
+
size = field.type.__size__
|
|
876
|
+
|
|
877
|
+
# Bit fields
|
|
878
|
+
if field.bits:
|
|
879
|
+
if self._starts_bit_value(field):
|
|
880
|
+
if not _can_pack_bits(field):
|
|
881
|
+
raise TypeError(f"Unsupported type for bit field: {field.type}")
|
|
882
|
+
|
|
883
|
+
yield from flush()
|
|
884
|
+
yield from align_to_field(field)
|
|
885
|
+
elif field.offset is not None:
|
|
886
|
+
raise TypeError(f"Unsupported offset for bit field: {field}")
|
|
887
|
+
else:
|
|
888
|
+
# Only the first bit field of a value moves the stream
|
|
889
|
+
size = 0
|
|
890
|
+
|
|
891
|
+
self._advance_bits(field)
|
|
892
|
+
current_bits.append(field)
|
|
893
|
+
|
|
894
|
+
# Basic and composite types (and arrays of them), and inlined structures
|
|
895
|
+
elif _can_pack(field.type):
|
|
896
|
+
if field.offset is None:
|
|
897
|
+
# The alignment depends on the position in the stream, so this needs a new block
|
|
898
|
+
if self.align:
|
|
899
|
+
yield from flush()
|
|
900
|
+
elif current_offset is None:
|
|
901
|
+
yield from flush()
|
|
902
|
+
elif current_block:
|
|
903
|
+
current_offset = field.offset
|
|
904
|
+
|
|
905
|
+
if not current_block:
|
|
906
|
+
yield from align_to_field(field)
|
|
907
|
+
current_block.append(field)
|
|
908
|
+
|
|
909
|
+
# Everything else - sub structures, arrays of structures, multi-dimensional and dynamic arrays
|
|
910
|
+
else:
|
|
911
|
+
yield from flush()
|
|
912
|
+
yield from align_to_field(field)
|
|
913
|
+
yield from self._generate_field_source(field)
|
|
914
|
+
|
|
915
|
+
# An aligned structure writes more than its size if the stream isn't aligned
|
|
916
|
+
if self.align:
|
|
917
|
+
size = None
|
|
918
|
+
|
|
919
|
+
if current_offset is not None:
|
|
920
|
+
current_offset = None if size is None else current_offset + size
|
|
921
|
+
|
|
922
|
+
yield from flush()
|
|
923
|
+
|
|
924
|
+
if self.align:
|
|
925
|
+
self._uses_start = True
|
|
926
|
+
yield 'stream.write(b"\\x00" * (-stream.tell() & (cls.__alignment__ - 1)))'
|
|
927
|
+
|
|
928
|
+
yield "return stream.tell() - start" if self._uses_start else f"return {current_offset}"
|
|
929
|
+
|
|
930
|
+
def _generate_field_source(self, field: Field) -> Iterator[str]:
|
|
931
|
+
"""Generate the code for writing a field using its type."""
|
|
932
|
+
self._uses_start = True
|
|
933
|
+
type_ = self._type_symbol(field.type)
|
|
934
|
+
lines = [
|
|
935
|
+
f"_value = {_get_attribute('data', field._name)}",
|
|
936
|
+
f"if _value is None: _value = {type_}.__default__()",
|
|
937
|
+
f"{type_}._write(stream, _value, endian=endian)",
|
|
938
|
+
]
|
|
939
|
+
|
|
940
|
+
yield _block(lines)
|
|
941
|
+
|
|
942
|
+
def _generate_bits_source(self, fields: list[Field]) -> Iterator[str]:
|
|
943
|
+
"""Generate the code for writing the bit fields that share a value of the bit buffer."""
|
|
944
|
+
type_ = _get_bits_type(fields[0])
|
|
945
|
+
|
|
946
|
+
pack = _Pack()
|
|
947
|
+
little = []
|
|
948
|
+
other = []
|
|
949
|
+
remaining = type_.__size__ * 8
|
|
950
|
+
for field in fields:
|
|
951
|
+
value = self._load_value(field, _get_attribute("data", field._name), pack)
|
|
952
|
+
if issubclass(field.type, (Enum, Flag)):
|
|
953
|
+
value += "._value_"
|
|
954
|
+
|
|
955
|
+
remaining -= field.bits
|
|
956
|
+
little.append(f"{value} << {type_.__size__ * 8 - remaining - field.bits}")
|
|
957
|
+
other.append(f"{value} << {remaining}")
|
|
958
|
+
|
|
959
|
+
# The bit buffer only writes the least significant bits first for explicit little endian
|
|
960
|
+
self._pack_number(type_, f"({' | '.join(little)} if endian == '<' else {' | '.join(other)})", pack)
|
|
961
|
+
|
|
962
|
+
yield self._generate_pack_source(pack)
|
|
963
|
+
|
|
964
|
+
def _generate_packed_source(self, fields: list[Field]) -> Iterator[str]:
|
|
965
|
+
"""Generate the code for writing consecutive packed fields."""
|
|
966
|
+
pack = _Pack()
|
|
967
|
+
self._pack_fields(fields, "data", pack)
|
|
968
|
+
|
|
969
|
+
if pack.items:
|
|
970
|
+
yield self._generate_pack_source(pack)
|
|
971
|
+
|
|
972
|
+
def _generate_pack_source(self, pack: _Pack) -> str:
|
|
973
|
+
"""Generate the code for checking, packing and writing values."""
|
|
974
|
+
function = f"{self._add_symbol(pack.functions('pack'), 'pack')}[endian]"
|
|
975
|
+
if pack.values == pack.retry:
|
|
976
|
+
return _block([*pack.lines, f"stream.write({function}({', '.join(pack.values)}))"])
|
|
977
|
+
|
|
978
|
+
# Numbers aren't checked for None up front, since struct.pack raises an error for those
|
|
979
|
+
return _block(
|
|
980
|
+
[
|
|
981
|
+
*pack.lines,
|
|
982
|
+
"try:",
|
|
983
|
+
f" buf = {function}({', '.join(pack.values)})",
|
|
984
|
+
"except _struct_error:",
|
|
985
|
+
f" buf = {function}({', '.join(pack.retry)})",
|
|
986
|
+
"stream.write(buf)",
|
|
987
|
+
]
|
|
988
|
+
)
|
|
989
|
+
|
|
990
|
+
def _load_value(self, field: Field, value: str, pack: _Pack) -> str:
|
|
991
|
+
"""Generate the code for loading a value, using the default of its type if it's ``None``."""
|
|
992
|
+
var = self._var("value")
|
|
993
|
+
pack.lines.append(f"if ({var} := {value}) is None: {var} = {self._type_symbol(field.type)}.__default__()")
|
|
994
|
+
return var
|
|
995
|
+
|
|
996
|
+
def _pack_fields(self, fields: list[Field], var: str, pack: _Pack, proxied: bool = False) -> None:
|
|
997
|
+
"""Generate the code for packing the given fields of the object in ``var``.
|
|
998
|
+
|
|
999
|
+
The fields of objects in a union are proxied, in which case we use the object behind the proxy.
|
|
1000
|
+
"""
|
|
1001
|
+
for field, count, char in _generate_struct_info(fields, inline=_inline_kind):
|
|
1002
|
+
if field is None:
|
|
1003
|
+
# Padding
|
|
1004
|
+
pack.add(count, "x")
|
|
1005
|
+
continue
|
|
1006
|
+
|
|
1007
|
+
value = _get_attribute(var, field._name)
|
|
1008
|
+
|
|
1009
|
+
if char in ("S", "U"):
|
|
1010
|
+
# Inlined structure or union
|
|
1011
|
+
self._witness(field.type)
|
|
1012
|
+
nested = self._load_value(field, value, pack)
|
|
1013
|
+
if proxied:
|
|
1014
|
+
pack.lines.append(f"if type({nested}) is _UnionProxy: {nested} = {nested}.__target__")
|
|
1015
|
+
|
|
1016
|
+
if char == "S":
|
|
1017
|
+
self._pack_fields(field.type.__fields__, nested, pack, proxied)
|
|
1018
|
+
else:
|
|
1019
|
+
self._pack_union(nested, field.type, pack)
|
|
1020
|
+
continue
|
|
1021
|
+
|
|
1022
|
+
self._pack_field(field, value, count, char, pack)
|
|
1023
|
+
|
|
1024
|
+
def _pack_union(self, var: str, type_: type[Union], pack: _Pack) -> None:
|
|
1025
|
+
"""Generate the code for packing the union in ``var``, by the field the interpreted writer uses."""
|
|
1026
|
+
# Changing a nested structure can change which field is the largest
|
|
1027
|
+
for field in type_.__fields__:
|
|
1028
|
+
if issubclass(field.type, Structure):
|
|
1029
|
+
self._witness(field.type)
|
|
1030
|
+
|
|
1031
|
+
self._pack_fields([type_._write_field()], var, pack, proxied=True)
|
|
1032
|
+
|
|
1033
|
+
def _pack_field(self, field: Field, value: str, count: int, char: str, pack: _Pack) -> None:
|
|
1034
|
+
"""Generate the code for checking and packing a field."""
|
|
1035
|
+
field_type = field.type
|
|
1036
|
+
read_type = _get_read_type(field_type)
|
|
1037
|
+
type_ = self._type_symbol(field_type)
|
|
1038
|
+
|
|
1039
|
+
if issubclass(field_type, Array):
|
|
1040
|
+
read_type = _get_read_type(field_type.__type__)
|
|
1041
|
+
var = self._var("value")
|
|
1042
|
+
num = field_type.__num_entries__
|
|
1043
|
+
pack.lines.append(
|
|
1044
|
+
f"if ({var} := {value}) is None or len({var}) != {num}: {var} = _array_values({type_}, {var})"
|
|
1045
|
+
)
|
|
1046
|
+
|
|
1047
|
+
# Entries can be enum members or plain integers, both can be packed as they are
|
|
1048
|
+
if issubclass(read_type, Int):
|
|
1049
|
+
pack.add(count, "s", f"b''.join([{_int_to_bytes(read_type, 'item')} for item in {var}])")
|
|
1050
|
+
else:
|
|
1051
|
+
pack.add(count, char, f"*{var}")
|
|
1052
|
+
|
|
1053
|
+
elif issubclass(field_type, (Char, CharArray, Wchar, WcharArray)):
|
|
1054
|
+
var = self._var("value")
|
|
1055
|
+
|
|
1056
|
+
# The format pads values that are too short with zeroes
|
|
1057
|
+
if issubclass(field_type, (Wchar, WcharArray)):
|
|
1058
|
+
encoding = self._add_symbol(Wchar.__encoding_map__, "encoding")
|
|
1059
|
+
pack.lines.append(
|
|
1060
|
+
f"if ({var} := {value}) is None or "
|
|
1061
|
+
f"len({var} := {var}.encode({encoding}[endian], 'surrogatepass')) > {count}: "
|
|
1062
|
+
f"{var} = _field_bytes({type_}, {value}, endian)"
|
|
1063
|
+
)
|
|
1064
|
+
else:
|
|
1065
|
+
pack.lines.append(
|
|
1066
|
+
f"if not isinstance({var} := {value}, bytes) or len({var}) > {count}: "
|
|
1067
|
+
f"{var} = _field_bytes({type_}, {var}, endian)"
|
|
1068
|
+
)
|
|
1069
|
+
|
|
1070
|
+
pack.add(count, "s", var)
|
|
1071
|
+
|
|
1072
|
+
elif issubclass(field_type, (Enum, Flag)):
|
|
1073
|
+
self._pack_number(read_type, f"{self._load_value(field, value, pack)}._value_", pack)
|
|
1074
|
+
|
|
1075
|
+
elif issubclass(read_type, Packed):
|
|
1076
|
+
pack.add(1, read_type.__fmt__, value, f"_or_default({value}, {type_})")
|
|
1077
|
+
|
|
1078
|
+
else:
|
|
1079
|
+
self._pack_number(read_type, self._load_value(field, value, pack), pack)
|
|
1080
|
+
|
|
1081
|
+
def _pack_number(self, type_: type[BaseType], value: str, pack: _Pack) -> None:
|
|
1082
|
+
"""Generate the code for packing a value of an integer or packed type."""
|
|
1083
|
+
if issubclass(type_, Int):
|
|
1084
|
+
pack.add(type_.__size__, "s", _int_to_bytes(type_, value))
|
|
1085
|
+
else:
|
|
1086
|
+
pack.add(1, type_.__fmt__, value)
|
|
1087
|
+
|
|
1088
|
+
|
|
1089
|
+
class _Pack:
|
|
1090
|
+
"""Tracks consecutive struct pack formats, and the code for the values to pack."""
|
|
1091
|
+
|
|
1092
|
+
def __init__(self):
|
|
1093
|
+
self.items: list[tuple[int, str]] = []
|
|
1094
|
+
self.lines: list[str] = []
|
|
1095
|
+
self.values: list[str] = []
|
|
1096
|
+
self.retry: list[str] = []
|
|
1097
|
+
|
|
1098
|
+
def add(self, count: int, char: str, value: str | None = None, retry: str | None = None) -> None:
|
|
1099
|
+
self.items.append((count, char))
|
|
1100
|
+
if value is not None:
|
|
1101
|
+
self.values.append(value)
|
|
1102
|
+
self.retry.append(retry or value)
|
|
1103
|
+
|
|
1104
|
+
def functions(self, name: str) -> dict[str, Callable]:
|
|
1105
|
+
"""Return the given function of the struct for every endianness."""
|
|
1106
|
+
return _struct_functions(_optimize_struct_fmt(self.items), name)
|
|
1107
|
+
|
|
1108
|
+
|
|
730
1109
|
def _check_union(cls: type[Union], fields: list[Field], inline: bool = True) -> int:
|
|
731
1110
|
"""Check if the compiler supports the union. Returns its size."""
|
|
732
1111
|
size = cls.__size__
|
|
@@ -796,6 +1175,52 @@ def _uses_generated_init(type_: type[Structure]) -> bool:
|
|
|
796
1175
|
return "__init__" in type_.__dict__ and type_.__new__ is object.__new__
|
|
797
1176
|
|
|
798
1177
|
|
|
1178
|
+
def _can_pack(type_: type[BaseType]) -> bool:
|
|
1179
|
+
"""Determine if a value of a type can be written by packing it, instead of by using the ``_write`` method."""
|
|
1180
|
+
if type_.__size__ is None:
|
|
1181
|
+
return False
|
|
1182
|
+
|
|
1183
|
+
if issubclass(type_, (Enum, Flag)):
|
|
1184
|
+
return _can_pack(type_.__type__)
|
|
1185
|
+
|
|
1186
|
+
if issubclass(type_, Structure):
|
|
1187
|
+
kind = _inline_kind(type_)
|
|
1188
|
+
if kind == "U":
|
|
1189
|
+
field = type_._write_field()
|
|
1190
|
+
return field is not None and bool(field.type.__size__) and _can_pack(field.type)
|
|
1191
|
+
return kind == "S" and all(_can_pack(field.type) for field in type_.__fields__)
|
|
1192
|
+
|
|
1193
|
+
if issubclass(type_, Array) and (
|
|
1194
|
+
issubclass(type_.__type__, (Structure, BaseArray)) or not _can_pack(type_.__type__)
|
|
1195
|
+
):
|
|
1196
|
+
return False
|
|
1197
|
+
|
|
1198
|
+
return any(issubclass(type_, base) and _has_default_write(type_, base) for base in _PACKABLE_TYPES)
|
|
1199
|
+
|
|
1200
|
+
|
|
1201
|
+
def _can_pack_bits(field: Field) -> bool:
|
|
1202
|
+
"""Determine if the value of the bit buffer for a bit field can be written by packing it."""
|
|
1203
|
+
type_ = _get_bits_type(field)
|
|
1204
|
+
return issubclass(type_, (Int, Packed)) and _can_pack(type_)
|
|
1205
|
+
|
|
1206
|
+
|
|
1207
|
+
def _has_default_write(type_: type[BaseType], base: type[BaseType]) -> bool:
|
|
1208
|
+
"""Return whether a type uses the write methods of its base type, and doesn't have custom ones."""
|
|
1209
|
+
return all(getattr(type_, name).__func__ is getattr(base, name).__func__ for name in ("_write", "_write_array"))
|
|
1210
|
+
|
|
1211
|
+
|
|
1212
|
+
def _int_to_bytes(type_: type[Int], value: str) -> str:
|
|
1213
|
+
"""Return the code to convert a value of an integer type to bytes."""
|
|
1214
|
+
return f"{value}.to_bytes({type_.__size__}, _byteorder[endian], signed={bool(type_.signed)})"
|
|
1215
|
+
|
|
1216
|
+
|
|
1217
|
+
def _get_attribute(var: str, name: str) -> str:
|
|
1218
|
+
"""Return the code to get an attribute."""
|
|
1219
|
+
if name.isidentifier() and not iskeyword(name):
|
|
1220
|
+
return f"{var}.{name}"
|
|
1221
|
+
return f"_getattr({var}, {name!r})"
|
|
1222
|
+
|
|
1223
|
+
|
|
799
1224
|
def _set_attribute(var: str, name: str, value: str) -> str:
|
|
800
1225
|
"""Return the code to set an attribute."""
|
|
801
1226
|
if name.isidentifier() and not iskeyword(name):
|
|
@@ -384,12 +384,24 @@ class BaseArray(BaseType):
|
|
|
384
384
|
if cls.__null_terminated__:
|
|
385
385
|
return cls.__type__._write_0(stream, data, endian=endian)
|
|
386
386
|
|
|
387
|
-
if not cls.__dynamic__
|
|
387
|
+
if not cls.__dynamic__:
|
|
388
|
+
data = cls._pad(data)
|
|
389
|
+
|
|
390
|
+
return cls.__type__._write_array(stream, data, endian=endian)
|
|
391
|
+
|
|
392
|
+
@classmethod
|
|
393
|
+
def _pad(cls, data: list[Any]) -> list[Any]:
|
|
394
|
+
"""Internal function for padding the entries of a static array to its size.
|
|
395
|
+
|
|
396
|
+
Args:
|
|
397
|
+
data: The entries of the array.
|
|
398
|
+
"""
|
|
399
|
+
if remaining := cls.__num_entries__ - (actual_size := len(data)):
|
|
388
400
|
if remaining < 0:
|
|
389
401
|
raise ArraySizeError(f"Expected static array size {cls.__num_entries__}, got {actual_size} instead.")
|
|
390
402
|
data = data + [cls.__type__.__default__()] * remaining
|
|
391
403
|
|
|
392
|
-
return
|
|
404
|
+
return data
|
|
393
405
|
|
|
394
406
|
@classproperty
|
|
395
407
|
def type(cls) -> type[BaseType]:
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: dissect.cstruct
|
|
3
|
-
Version: 5.0.
|
|
3
|
+
Version: 5.0.dev14
|
|
4
4
|
Summary: A Dissect module implementing a parser for C-like structures: structure parsing in Python made easy
|
|
5
5
|
Author-email: Dissect Team <dissect@fox-it.com>
|
|
6
6
|
License-Expression: Apache-2.0
|
|
@@ -404,3 +404,55 @@ def test_align_after_dynamic_field(cs: cstruct, compiled: bool) -> None:
|
|
|
404
404
|
assert obj.c == 0x04030201
|
|
405
405
|
assert obj.d == 0xCC
|
|
406
406
|
assert obj.dumps() == buf
|
|
407
|
+
|
|
408
|
+
|
|
409
|
+
def test_align_stream_position(cs: cstruct, compiled: bool) -> None:
|
|
410
|
+
"""Test that padding after a dynamic field or at the end of a structure follows the position in the stream."""
|
|
411
|
+
cdef = """
|
|
412
|
+
struct test_static {
|
|
413
|
+
uint8 a;
|
|
414
|
+
uint32 b;
|
|
415
|
+
uint8 c;
|
|
416
|
+
};
|
|
417
|
+
|
|
418
|
+
struct test_dynamic {
|
|
419
|
+
uint8 n;
|
|
420
|
+
uint8 a[n];
|
|
421
|
+
uint32 b;
|
|
422
|
+
uint8 c;
|
|
423
|
+
};
|
|
424
|
+
|
|
425
|
+
struct nested {
|
|
426
|
+
uint8 a;
|
|
427
|
+
uint32 b;
|
|
428
|
+
};
|
|
429
|
+
|
|
430
|
+
struct test_nested {
|
|
431
|
+
uint8 c;
|
|
432
|
+
nested d;
|
|
433
|
+
uint64 e;
|
|
434
|
+
};
|
|
435
|
+
"""
|
|
436
|
+
cs.load(cdef, compiled=compiled, align=True)
|
|
437
|
+
|
|
438
|
+
assert verify_compiled(cs.test_static, compiled)
|
|
439
|
+
assert verify_compiled(cs.test_dynamic, compiled)
|
|
440
|
+
assert verify_compiled(cs.test_nested, compiled)
|
|
441
|
+
|
|
442
|
+
# The padding between static fields depends on the structure, the padding at the end on the stream
|
|
443
|
+
fh = BytesIO(b"\xff")
|
|
444
|
+
fh.seek(1)
|
|
445
|
+
assert cs.test_static(a=1, b=2, c=3).write(fh) == 11
|
|
446
|
+
assert fh.getvalue() == b"\xff\x01\x00\x00\x00\x02\x00\x00\x00\x03\x00\x00"
|
|
447
|
+
|
|
448
|
+
# The padding after a dynamic field depends on the stream
|
|
449
|
+
fh = BytesIO(b"\xff")
|
|
450
|
+
fh.seek(1)
|
|
451
|
+
assert cs.test_dynamic(n=1, a=[2], b=3, c=4).write(fh) == 11
|
|
452
|
+
assert fh.getvalue() == b"\xff\x01\x02\x00\x03\x00\x00\x00\x04\x00\x00\x00"
|
|
453
|
+
|
|
454
|
+
# The padding at the end of a nested structure depends on the stream too
|
|
455
|
+
fh = BytesIO(b"\xff")
|
|
456
|
+
fh.seek(1)
|
|
457
|
+
assert cs.test_nested(c=1, d=cs.nested(a=2, b=3), e=4).write(fh) == 31
|
|
458
|
+
assert fh.getvalue() == (b"\xff\x01\x00\x00\x00\x02\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x04" + bytes(14))
|