Allow enums to reference their own members

fox-it · Jul 28, 2023 · 2004de4 · 2004de4
1 parent 80b64cd
commit 2004de4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/dissect/cstruct/parser.py b/dissect/cstruct/parser.py
@@ -117,7 +117,7 @@ def _enum(self, tokens: TokenConsumer) -> None:
                 if not val:
                     val = nextval
                 else:
-                    val = Expression(self.cstruct, val).evaluate()
+                    val = Expression(self.cstruct, val).evaluate(values)
 
                 if enumtype == "flag":
                     high_bit = val.bit_length() - 1

diff --git a/tests/test_enum.py b/tests/test_enum.py
@@ -260,3 +260,19 @@ def test_enum_anonymous_struct(compiled):
 
     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]
+
+
+def test_enum_reference_own_member(compiled):
+    cdef = """
+    enum test {
+        A,
+        B = A + 3,
+        C
+    };
+    """
+    cs = cstruct.cstruct()
+    cs.load(cdef, compiled=compiled)
+
+    assert cs.test.A == 0
+    assert cs.test.B == 3
+    assert cs.test.C == 4