Skip to content

Commit

Permalink
Add the default() method back in
Browse files Browse the repository at this point in the history
  • Loading branch information
Schamper committed Nov 26, 2023
1 parent 5b38ad8 commit a30bde6
Show file tree
Hide file tree
Showing 5 changed files with 111 additions and 90 deletions.
7 changes: 7 additions & 0 deletions dissect/cstruct/types/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
8 changes: 8 additions & 0 deletions dissect/cstruct/types/char.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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))
2 changes: 1 addition & 1 deletion dissect/cstruct/types/structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions dissect/cstruct/types/wchar.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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))
176 changes: 87 additions & 89 deletions tests/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit a30bde6

Please sign in to comment.