From 57ae87a389052718bd4d3b0eba3e09994604d86f Mon Sep 17 00:00:00 2001 From: Schamper <1254028+Schamper@users.noreply.github.com> Date: Mon, 23 Sep 2024 14:18:30 +0200 Subject: [PATCH] Fix enum and flag repr in struct reprs --- dissect/cstruct/types/structure.py | 3 +- tests/test_types_enum.py | 86 ++++++++++++++++++++++-------- tests/test_types_flag.py | 52 ++++++++++++------ 3 files changed, 102 insertions(+), 39 deletions(-) diff --git a/dissect/cstruct/types/structure.py b/dissect/cstruct/types/structure.py index 817876d..84292cb 100644 --- a/dissect/cstruct/types/structure.py +++ b/dissect/cstruct/types/structure.py @@ -2,6 +2,7 @@ import io from contextlib import contextmanager +from enum import Enum from functools import lru_cache from operator import attrgetter from textwrap import dedent @@ -381,7 +382,7 @@ def __getitem__(self, item: str) -> Any: def __repr__(self) -> str: values = [ - f"{k}={hex(self[k]) if (issubclass(f.type, int) and not issubclass(f.type, Pointer)) else repr(self[k])}" + f"{k}={hex(self[k]) if (issubclass(f.type, int) and not issubclass(f.type, (Pointer, Enum))) else repr(self[k])}" for k, f in self.__class__.fields.items() ] return f"<{self.__class__.__name__} {' '.join(values)}>" diff --git a/tests/test_types_enum.py b/tests/test_types_enum.py index 95b3b3f..0a9c172 100644 --- a/tests/test_types_enum.py +++ b/tests/test_types_enum.py @@ -98,6 +98,66 @@ def test_enum_eof(TestEnum: type[Enum]) -> None: TestEnum[None](b"\x01") +def test_enum_same_value_different_type(cs: cstruct, compiled: bool) -> None: + cdef = """ + enum Test16 : uint16 { + A = 0x1, + B = 0x2 + }; + + enum Test24 : uint24 { + A = 0x1, + B = 0x2 + }; + + enum Test32 : uint32 { + A = 0x1, + B = 0x2 + }; + """ + cs.load(cdef, compiled=compiled) + + obj = { + cs.Test16.A: "Test16.A", + cs.Test16.B: "Test16.B", + cs.Test24.A: "Test24.A", + cs.Test24.B: "Test24.B", + } + + assert obj[cs.Test16.A] == "Test16.A" + assert obj[cs.Test16(2)] == "Test16.B" + assert obj[cs.Test24(1)] == "Test24.A" + assert obj[cs.Test24.B] == "Test24.B" + + with pytest.raises(KeyError): + obj[cs.Test32.A] + + +def test_enum_str_repr(TestEnum: type[Enum]) -> None: + assert repr(TestEnum.A) == "" + assert str(TestEnum.A) == "Test.A" + assert repr(TestEnum(69)) == "" + assert str(TestEnum(69)) == "Test.69" + + +def test_enum_str_repr_in_struct(cs: cstruct, compiled: bool) -> None: + cdef = """ + enum Test16 : uint16 { + A = 0x1, + B = 0x2 + }; + + struct test { + Test16 a; + }; + """ + cs.load(cdef, compiled=compiled) + + obj = cs.test(b"\x02\x00") + assert repr(obj) == ">" + assert str(obj) == ">" + + def test_enum_struct(cs: cstruct, compiled: bool) -> None: cdef = """ enum Test16 : uint16 { @@ -174,26 +234,6 @@ def test_enum_struct(cs: cstruct, compiled: bool) -> None: assert cs.test_expr(buf).expr == [cs.Test16.A, cs.Test16.B] assert cs.test_expr(size=1, expr=[cs.Test16.A, cs.Test16.B]).dumps() == buf - obj = { - cs.Test16.A: "Test16.A", - cs.Test16.B: "Test16.B", - cs.Test24.A: "Test24.A", - cs.Test24.B: "Test24.B", - } - - assert obj[cs.Test16.A] == "Test16.A" - assert obj[cs.Test16(2)] == "Test16.B" - assert obj[cs.Test24(1)] == "Test24.A" - assert obj[cs.Test24.B] == "Test24.B" - - with pytest.raises(KeyError): - obj[cs.Test32.A] - - assert repr(cs.Test16.A) == "" - assert str(cs.Test16.A) == "Test16.A" - assert repr(cs.Test16(69)) == "" - assert str(cs.Test16(69)) == "Test16.69" - def test_enum_comments(cs: cstruct) -> None: cdef = """ @@ -259,7 +299,7 @@ def test_enum_name(cs: cstruct, compiled: bool) -> None: Color = cs.Color Pixel = cs.Pixel - pixel = Pixel(b"\xFF\x0A\x01\x00\xAA\xBB\xCC\xDD") + pixel = Pixel(b"\xff\x0a\x01\x00\xaa\xbb\xcc\xdd") assert pixel.x == 255 assert pixel.y == 10 assert pixel.color.name == "RED" @@ -267,7 +307,7 @@ def test_enum_name(cs: cstruct, compiled: bool) -> None: assert pixel.color.value == 1 assert pixel.hue == 0xDDCCBBAA - pixel = Pixel(b"\x00\x00\xFF\x00\xAA\xBB\xCC\xDD") + pixel = Pixel(b"\x00\x00\xff\x00\xaa\xbb\xcc\xdd") assert pixel.color.name is None assert pixel.color.value == 0xFF assert repr(pixel.color) == "" @@ -350,7 +390,7 @@ def test_enum_anonymous_struct(cs: cstruct, compiled: bool) -> None: test = cs.test - t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0A\x00\x00\x00") + t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x00\x00") assert t.arr == [255, 0, 0, 10] diff --git a/tests/test_types_flag.py b/tests/test_types_flag.py index a326225..be78542 100644 --- a/tests/test_types_flag.py +++ b/tests/test_types_flag.py @@ -62,6 +62,42 @@ def test_flag_operator(TestFlag: type[Flag]) -> None: assert TestFlag[2]([TestFlag.B, (TestFlag.A | TestFlag.B)]).dumps() == b"\x02\x03" +def test_flag_str_repr(TestFlag: type[Flag]) -> None: + if PY_311: + assert repr(TestFlag.A | TestFlag.B) == "" + assert str(TestFlag.A | TestFlag.B) == "Test.A|B" + assert repr(TestFlag(69)) == "" + assert str(TestFlag(69)) == "Test.A|68" + else: + assert repr(TestFlag.A | TestFlag.B) == "" + assert str(TestFlag.A | TestFlag.B) == "Test.B|A" + assert repr(TestFlag(69)) == "" + assert str(TestFlag(69)) == "Test.68|A" + + +def test_flag_str_repr_in_struct(cs: cstruct, compiled: bool) -> None: + cdef = """ + flag Test : uint16 { + A, + B + }; + + struct test { + Test a; + }; + """ + cs.load(cdef, compiled=compiled) + + obj = cs.test(b"\x03\x00") + + if PY_311: + assert repr(obj) == ">" + assert str(obj) == ">" + else: + assert repr(obj) == ">" + assert str(obj) == ">" + + def test_flag_struct(cs: cstruct) -> None: cdef = """ flag Test { @@ -101,20 +137,6 @@ def test_flag_struct(cs: cstruct) -> None: assert bool(cs.Test(1)) is True assert cs.Test.a | cs.Test.b == 3 - if PY_311: - assert repr(cs.Test.c | cs.Test.d) == "" - assert str(cs.Test.c | cs.Test.d) == "Test.c|d" - assert repr(cs.Test.a | cs.Test.b) == "" - assert str(cs.Test.a | cs.Test.b) == "Test.a|b" - assert repr(cs.Test(69)) == "" - assert str(cs.Test(69)) == "Test.a|c|64" - else: - assert repr(cs.Test.c | cs.Test.d) == "" - assert str(cs.Test.c | cs.Test.d) == "Test.d|c" - assert repr(cs.Test.a | cs.Test.b) == "" - assert str(cs.Test.a | cs.Test.b) == "Test.b|a" - assert repr(cs.Test(69)) == "" - assert str(cs.Test(69)) == "Test.64|c|a" assert cs.Test(2) == cs.Test.b assert cs.Test(3) == cs.Test.a | cs.Test.b assert cs.Test.c & 12 == cs.Test.c @@ -231,5 +253,5 @@ def test_flag_anonymous_struct(cs: cstruct, compiled: bool) -> None: test = cs.test - t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0A\x00\x00\x00") + t = test(b"\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0a\x00\x00\x00") assert t.arr == [255, 0, 0, 10]