From 827a1d86a70ee29e6195b550e6fd89d29c1696bc Mon Sep 17 00:00:00 2001 From: Erik Schamper <1254028+Schamper@users.noreply.github.com> Date: Mon, 8 Jan 2024 17:35:46 +0100 Subject: [PATCH] Add the `default()` method back in (#55) --- dissect/cstruct/types/base.py | 7 ++ dissect/cstruct/types/char.py | 8 ++ dissect/cstruct/types/structure.py | 2 +- dissect/cstruct/types/wchar.py | 8 ++ tests/test_basic.py | 176 ++++++++++++++--------------- 5 files changed, 111 insertions(+), 90 deletions(-) diff --git a/dissect/cstruct/types/base.py b/dissect/cstruct/types/base.py index a7f5d6e..e55e138 100644 --- a/dissect/cstruct/types/base.py +++ b/dissect/cstruct/types/base.py @@ -54,6 +54,10 @@ def __len__(cls) -> int: return cls.size + def default(cls) -> BaseType: + """Return the default value of this type.""" + return cls() + def reads(cls, data: bytes) -> BaseType: """Parse the given data from a bytes-like object. @@ -196,6 +200,9 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Array: return cls.type._read_array(stream, num, context) + def default(cls) -> BaseType: + return [cls.type.default() for _ in range(0 if cls.dynamic or cls.null_terminated else cls.num_entries)] + class Array(list, BaseType, metaclass=ArrayMetaType): """Implements a fixed or dynamically sized array type. diff --git a/dissect/cstruct/types/char.py b/dissect/cstruct/types/char.py index ab84cad..8fe9436 100644 --- a/dissect/cstruct/types/char.py +++ b/dissect/cstruct/types/char.py @@ -48,6 +48,10 @@ def _write(cls, stream: BinaryIO, data: Union[bytes, int, str]) -> int: return stream.write(data) + @classmethod + def default(cls) -> Char: + return type.__call__(cls, b"\x00") + class CharArray(bytes, BaseType, metaclass=ArrayMetaType): """Character array type for reading and writing byte strings.""" @@ -67,3 +71,7 @@ def _write(cls, stream: BinaryIO, data: bytes) -> int: if cls.null_terminated: return stream.write(data + b"\x00") return stream.write(data) + + @classmethod + def default(cls) -> CharArray: + return type.__call__(cls, b"\x00" * (0 if cls.dynamic or cls.null_terminated else cls.num_entries)) diff --git a/dissect/cstruct/types/structure.py b/dissect/cstruct/types/structure.py index dba4f6c..5680d9e 100644 --- a/dissect/cstruct/types/structure.py +++ b/dissect/cstruct/types/structure.py @@ -63,7 +63,7 @@ def __call__(cls, *args, **kwargs) -> Structure: # Shortcut for single char/bytes type return type.__call__(cls, *args, **kwargs) elif not args and not kwargs: - obj = cls(**{field.name: field.type() for field in cls.__fields__}) + obj = cls(**{field.name: field.type.default() for field in cls.__fields__}) object.__setattr__(obj, "_values", {}) object.__setattr__(obj, "_sizes", {}) return obj diff --git a/dissect/cstruct/types/wchar.py b/dissect/cstruct/types/wchar.py index a823e8e..cd1a90c 100644 --- a/dissect/cstruct/types/wchar.py +++ b/dissect/cstruct/types/wchar.py @@ -54,6 +54,10 @@ def _read_0(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Wchar: def _write(cls, stream: BinaryIO, data: str) -> int: return stream.write(data.encode(cls.__encoding_map__[cls.cs.endian])) + @classmethod + def default(cls) -> Wchar: + return type.__call__(cls, "\x00") + class WcharArray(str, BaseType, metaclass=ArrayMetaType): """Wide-character array type for reading and writing UTF-16 strings.""" @@ -67,3 +71,7 @@ def _write(cls, stream: BinaryIO, data: str) -> int: if cls.null_terminated: data += "\x00" return stream.write(data.encode(Wchar.__encoding_map__[cls.cs.endian])) + + @classmethod + def default(cls) -> WcharArray: + return type.__call__(cls, "\x00" * (0 if cls.dynamic or cls.null_terminated else cls.num_entries)) diff --git a/tests/test_basic.py b/tests/test_basic.py index 8717b45..1316c25 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -107,95 +107,93 @@ def test_lookups(cs: cstruct, compiled: bool): assert cs.lookups["a"] == {1: 3, 2: 4} -# TODO: -# def test_default_constructors(cs: cstruct, compiled: bool): -# cdef = """ -# enum Enum { -# a = 0, -# b = 1 -# }; - -# flag Flag { -# a = 0, -# b = 1 -# }; - -# struct test { -# uint32 t_int; -# uint32 t_int_array[2]; -# uint24 t_bytesint; -# uint24 t_bytesint_array[2]; -# char t_char; -# char t_char_array[2]; -# wchar t_wchar; -# wchar t_wchar_array[2]; -# Enum t_enum; -# Enum t_enum_array[2]; -# Flag t_flag; -# Flag t_flag_array[2]; -# }; -# """ -# cs.load(cdef, compiled=compiled) - -# assert verify_compiled(cs.test, compiled) - -# obj = cs.test() -# assert obj.t_int == 0 -# assert obj.t_int_array == [0, 0] -# assert obj.t_bytesint == 0 -# assert obj.t_bytesint_array == [0, 0] -# assert obj.t_char == b"\x00" -# assert obj.t_char_array == b"\x00\x00" -# assert obj.t_wchar == "\x00" -# assert obj.t_wchar_array == "\x00\x00" -# assert obj.t_enum == cs.Enum(0) -# assert obj.t_enum_array == [cs.Enum(0), cs.Enum(0)] -# assert obj.t_flag == cs.Flag(0) -# assert obj.t_flag_array == [cs.Flag(0), cs.Flag(0)] - -# assert obj.dumps() == b"\x00" * 54 - - -# TODO: -# def test_default_constructors_dynamic(cs: cstruct, compiled: bool): -# cdef = """ -# enum Enum { -# a = 0, -# b = 1 -# }; -# flag Flag { -# a = 0, -# b = 1 -# }; -# struct test { -# uint8 x; -# uint32 t_int_array_n[]; -# uint32 t_int_array_d[x]; -# uint24 t_bytesint_array_n[]; -# uint24 t_bytesint_array_d[x]; -# char t_char_array_n[]; -# char t_char_array_d[x]; -# wchar t_wchar_array_n[]; -# wchar t_wchar_array_d[x]; -# Enum t_enum_array_n[]; -# Enum t_enum_array_d[x]; -# Flag t_flag_array_n[]; -# Flag t_flag_array_d[x]; -# }; -# """ -# cs.load(cdef, compiled=compiled) - -# assert verify_compiled(cs.test, compiled) - -# obj = cs.test() - -# assert obj.t_int_array_n == obj.t_int_array_d == [] -# assert obj.t_bytesint_array_n == obj.t_bytesint_array_d == [] -# assert obj.t_char_array_n == obj.t_char_array_d == b"" -# assert obj.t_wchar_array_n == obj.t_wchar_array_d == "" -# assert obj.t_enum_array_n == obj.t_enum_array_d == [] -# assert obj.t_flag_array_n == obj.t_flag_array_d == [] -# assert obj.dumps() == b"\x00" * 19 +def test_default_constructors(cs: cstruct, compiled: bool): + cdef = """ + enum Enum { + a = 0, + b = 1 + }; + + flag Flag { + a = 0, + b = 1 + }; + + struct test { + uint32 t_int; + uint32 t_int_array[2]; + uint24 t_bytesint; + uint24 t_bytesint_array[2]; + char t_char; + char t_char_array[2]; + wchar t_wchar; + wchar t_wchar_array[2]; + Enum t_enum; + Enum t_enum_array[2]; + Flag t_flag; + Flag t_flag_array[2]; + }; + """ + cs.load(cdef, compiled=compiled) + + assert verify_compiled(cs.test, compiled) + + obj = cs.test() + assert obj.t_int == 0 + assert obj.t_int_array == [0, 0] + assert obj.t_bytesint == 0 + assert obj.t_bytesint_array == [0, 0] + assert obj.t_char == b"\x00" + assert obj.t_char_array == b"\x00\x00" + assert obj.t_wchar == "\x00" + assert obj.t_wchar_array == "\x00\x00" + assert obj.t_enum == cs.Enum(0) + assert obj.t_enum_array == [cs.Enum(0), cs.Enum(0)] + assert obj.t_flag == cs.Flag(0) + assert obj.t_flag_array == [cs.Flag(0), cs.Flag(0)] + + assert obj.dumps() == b"\x00" * 54 + + +def test_default_constructors_dynamic(cs: cstruct, compiled: bool): + cdef = """ + enum Enum { + a = 0, + b = 1 + }; + flag Flag { + a = 0, + b = 1 + }; + struct test { + uint8 x; + uint32 t_int_array_n[]; + uint32 t_int_array_d[x]; + uint24 t_bytesint_array_n[]; + uint24 t_bytesint_array_d[x]; + char t_char_array_n[]; + char t_char_array_d[x]; + wchar t_wchar_array_n[]; + wchar t_wchar_array_d[x]; + Enum t_enum_array_n[]; + Enum t_enum_array_d[x]; + Flag t_flag_array_n[]; + Flag t_flag_array_d[x]; + }; + """ + cs.load(cdef, compiled=compiled) + + assert verify_compiled(cs.test, compiled) + + obj = cs.test() + + assert obj.t_int_array_n == obj.t_int_array_d == [] + assert obj.t_bytesint_array_n == obj.t_bytesint_array_d == [] + assert obj.t_char_array_n == obj.t_char_array_d == b"" + assert obj.t_wchar_array_n == obj.t_wchar_array_d == "" + assert obj.t_enum_array_n == obj.t_enum_array_d == [] + assert obj.t_flag_array_n == obj.t_flag_array_d == [] + assert obj.dumps() == b"\x00" * 19 def test_config_flag_nocompile(cs: cstruct, compiled: bool):