Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions dissect/cstruct/bitbuffer.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,11 @@ def __init__(self, stream: BinaryIO, *, endian: str):

def read(self, field_type: type[BaseType], bits: int) -> int:
if self._remaining == 0 or self._type != field_type:
if field_type.size is None:
if field_type.__size__ is None:
raise ValueError("Reading variable-length fields is unsupported")

self._type = field_type
self._remaining = field_type.size * 8
self._remaining = field_type.__size__ * 8
self._buffer = field_type._read(self.stream, endian=self.endian)

if isinstance(self._buffer, bytes):
Expand Down Expand Up @@ -51,17 +51,17 @@ def write(self, field_type: type[BaseType], data: int, bits: int) -> None:
if self._type:
self.flush()

if field_type.size is None:
if field_type.__size__ is None:
raise ValueError("Writing variable-length fields is unsupported")

self._remaining = field_type.size * 8
self._remaining = field_type.__size__ * 8
self._type = field_type

if self._type is None or self._type.size is None:
if self._type is None or self._type.__size__ is None:
raise ValueError("Invalid state")

if self.endian == "<":
self._buffer |= data << (self._type.size * 8 - self._remaining)
self._buffer |= data << (self._type.__size__ * 8 - self._remaining)
else:
self._buffer |= data << (self._remaining - bits)

Expand Down
32 changes: 16 additions & 16 deletions dissect/cstruct/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@


def compile(structure: type[Structure]) -> type[Structure]:
return Compiler(structure.cs).compile(structure)
return Compiler(structure.__cs__).compile(structure)


class Compiler:
Expand Down Expand Up @@ -119,7 +119,7 @@ def generate_source(self) -> str:
"""

if any(field.bits for field in self.fields):
preamble += "bit_reader = BitBuffer(stream, endian=endian)\n"
preamble += "bit_reader = BitBuffer(stream, endian=cls.__cs__.endian)\n"

read_code = "\n".join(self._generate_fields())

Expand Down Expand Up @@ -224,22 +224,22 @@ def align_to_field(field: Field) -> Iterator[str]:
yield from flush()

if self.align:
yield f"stream.seek(-stream.tell() & (cls.alignment - 1), {io.SEEK_CUR})"
yield f"stream.seek(-stream.tell() & (cls.__alignment__ - 1), {io.SEEK_CUR})"

def _generate_structure(self, field: Field) -> Iterator[str]:
template = f"""
{"_s = stream.tell()" if field.type.dynamic else ""}
{"_s = stream.tell()" if field.type.__dynamic__ else ""}
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 ""}
{f's["{field._name}"] = stream.tell() - _s' if field.type.__dynamic__ else ""}
"""

yield dedent(template)

def _generate_array(self, field: Field) -> Iterator[str]:
template = f"""
{"_s = stream.tell()" if field.type.dynamic else ""}
{"_s = stream.tell()" if field.type.__dynamic__ else ""}
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 ""}
{f's["{field._name}"] = stream.tell() - _s' if field.type.__dynamic__ else ""}
"""

yield dedent(template)
Expand All @@ -253,8 +253,8 @@ def _generate_bits(self, field: Field) -> Iterator[str]:
field_type = field_type.type

if issubclass(field_type, Char):
field_type = field_type.cs.uint8
lookup = "cls.cs.uint8"
field_type = field_type.__cs__.uint8
lookup = "cls.__cs__.uint8"

template = f"""
_t = {lookup}
Expand Down Expand Up @@ -283,13 +283,13 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
read_type = _get_read_type(self.cs, field_type.type)

if issubclass(read_type, (Char, Wchar, Int)):
count *= read_type.size
count *= read_type.__size__
getter = f"buf[{size}:{size + count}]"
else:
getter = f"data[{slice_index}:{slice_index + count}]"
slice_index += count
elif issubclass(read_type, (Char, Wchar, Int)):
getter = f"buf[{size}:{size + read_type.size}]"
getter = f"buf[{size}:{size + read_type.__size__}]"
else:
getter = f"data[{slice_index}]"
slice_index += 1
Expand All @@ -308,8 +308,8 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:

if issubclass(field_type.type, Int):
reads.append(f"_b = {getter}")
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})]"
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, context=r, endian=endian)"
list_comp = f"[{item_parser} for e in {getter}]"
Expand All @@ -329,7 +329,7 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
reads.append(f'r["{field._name}"] = {parser}')
reads.append("") # Generates a newline in the resulting code

size += field_type.size
size += field_type.__size__

fmt = _optimize_struct_fmt(info)
if fmt == "x" or (len(fmt) == 2 and fmt[1] == "x"):
Expand Down Expand Up @@ -382,9 +382,9 @@ def _generate_struct_info(cs: cstruct, fields: list[Field], align: bool = False)
# Other types are byte based
# We don't actually unpack anything here but slice directly out of the buffer
elif issubclass(read_type, (Char, Wchar, Int)):
yield field, count * read_type.size, "x"
yield field, count * read_type.__size__, "x"

size = count * read_type.size
size = count * read_type.__size__
imaginary_offset += size
if current_offset is not None:
current_offset += size
Expand Down
54 changes: 42 additions & 12 deletions dissect/cstruct/cstruct.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,7 +206,28 @@ def __init__(self, load: str = "", *, endian: AllowedEndianness = "<", pointer:
if load:
self.load(load)

def __getattr__(self, attr: str) -> Any:
def __getattr__(
self, attr: str
) -> (
type[
LEB128
| BaseType
| Char
| Enum
| Flag
| Int
| Packed[int]
| Packed[float]
| Pointer
| Structure
| Union
| Void
| Wchar
]
| int
| str
| bytes
):
try:
return self.consts[attr]
except KeyError:
Expand All @@ -227,7 +248,7 @@ def __copy__(self) -> cstruct:
# 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
new_type.__cs__ = cs
cs.add_type(name, new_type, replace=True)

for name, value in self.consts.items():
Expand Down Expand Up @@ -481,10 +502,10 @@ def _make_type(
attrs = attrs or {}
attrs.update(
{
"cs": self,
"size": size,
"dynamic": size is None,
"alignment": alignment or size,
"__cs__": self,
"__size__": size,
"__dynamic__": size is None,
"__alignment__": alignment or size,
}
)
return types.new_class(name, bases, {}, lambda ns: ns.update(attrs))
Expand All @@ -494,12 +515,12 @@ def _make_array(self, type_: T, num_entries: int | Expression | None) -> type[Ar
if num_entries is None:
null_terminated = True
size = None
elif isinstance(num_entries, Expression) or type_.dynamic:
elif isinstance(num_entries, Expression) or type_.__dynamic__:
size = None
else:
if type_.size is None:
if type_.__size__ is None:
raise ValueError(f"Cannot create array of dynamic type: {type_.__name__}")
size = num_entries * type_.size
size = num_entries * type_.__size__

name = f"{type_.__name__}[]" if null_terminated else f"{type_.__name__}[{num_entries}]"

Expand All @@ -511,7 +532,7 @@ def _make_array(self, type_: T, num_entries: int | Expression | None) -> type[Ar
"null_terminated": null_terminated,
}

return cast("type[Array]", self._make_type(name, bases, size, alignment=type_.alignment, attrs=attrs))
return cast("type[Array]", self._make_type(name, bases, size, alignment=type_.__alignment__, attrs=attrs))

def _make_int_type(self, name: str, size: int, signed: bool, *, alignment: int | None = None) -> type[Int]:
return cast("type[Int]", self._make_type(name, (Int,), size, alignment=alignment, attrs={"signed": signed}))
Expand All @@ -538,8 +559,8 @@ def _make_pointer(self, target: type[BaseType]) -> type[Pointer]:
return self._make_type(
f"{target.__name__}*",
(Pointer,),
self.pointer.size,
alignment=self.pointer.alignment,
self.pointer.__size__,
alignment=self.pointer.__alignment__,
attrs={"type": target},
)

Expand Down Expand Up @@ -568,6 +589,15 @@ def _make_union(
) -> type[Structure]:
return self._make_struct(name, fields, align=align, anonymous=anonymous, base=Union)

@property
def typedefs(self) -> dict[str, type[BaseType]]:
warnings.warn(
"The 'typedefs' property is deprecated, use 'types' instead.",
DeprecationWarning,
stacklevel=2,
)
return self.types

if TYPE_CHECKING:
# ruff: noqa: PYI042
_int = int
Expand Down
4 changes: 2 additions & 2 deletions dissect/cstruct/tools/stubgen.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,15 +159,15 @@ def generate_structure_stub(
indent = " " * 4

args = ["self"]
for field_name, field in structure.fields.items():
for field_name, field in structure.__members__.items():
inlined = False

# If it's a structure and not globally defined, add an inline stub for it
nested_type = field.type
while issubclass(nested_type, types.BaseArray):
nested_type = nested_type.type

if issubclass(nested_type, types.Structure) and nested_type.__name__ not in structure.cs.types:
if issubclass(nested_type, types.Structure) and nested_type.__name__ not in structure.__cs__.types:
inlined = True
inline_stub = generate_structure_stub(nested_type, cs_prefix=cs_prefix, module_prefix=module_prefix)

Expand Down
Loading
Loading