From b68caf0375a718e3da0b5e87c7e33c2900781527 Mon Sep 17 00:00:00 2001
From: Cervenka Dusan <cervenka@acrios.com>
Date: Wed, 6 Apr 2022 02:38:05 +0200
Subject: [PATCH] Added support for simple macros

Signed-off-by: Cervenka Dusan <cervenka@acrios.com>
---
 dissect/cstruct/cstruct.py | 15 ++++++++++++---
 tests/test_basic.py        |  4 +++-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/dissect/cstruct/cstruct.py b/dissect/cstruct/cstruct.py
index a766d56..3a45d88 100644
--- a/dissect/cstruct/cstruct.py
+++ b/dissect/cstruct/cstruct.py
@@ -324,7 +324,7 @@ def parse(self, data):
         self._constants(data)
         self._enums(data)
         self._structs(data)
-        self._lookups(data, self.cstruct.consts)
+        self._lookups(data)
 
     def _constants(self, data):
         r = re.finditer(r'#define\s+(?P<name>[^\s]+)\s+(?P<value>[^\r\n]+)\s*\n', data)
@@ -375,7 +375,7 @@ def _parse_fields_enums(self, s):
             commentAttributes = self.parse_comment_block(d['commentBlock'])
 
             field = re.finditer(
-                r'(?P<key>[a-zA-z][^ =]*)[ ]*=?[ ]*(?P<value>[^ ]+)?',
+                r'(?P<key>[a-zA-z][^ =]*)[ ]*=?[ ]*(?P<value>[^,/]+)?',
                 d["value"],
             )
 
@@ -484,7 +484,7 @@ def parse_comment_block(self,s):
 
         return commentAttributes
 
-    def _lookups(self, data, consts):
+    def _lookups(self, data):
         r = re.finditer(r'\$(?P<name>[^\s]+) = ({[^}]+})\w*\n', data)
 
         for t in r:
@@ -493,6 +493,15 @@ def _lookups(self, data, consts):
                 [(self.cstruct.consts[k], v) for k, v in d.items()]
             )
 
+        #support simple defines
+        r = re.finditer(r'^[ \t]*#define (?P<name>[^\s]+)[ \t]+(?P<value>\(?-?[0-9]+U?\)?)\n?', data)
+
+        for t in r:
+            d = ast.literal_eval(t.group(2))
+            self.cstruct.lookups[t.group(1)] = dict(
+                [(self.cstruct.consts[k], v) for k, v in d.items()]
+            )
+
 
 class Instance(object):
     """Holds parsed structure data."""
diff --git a/tests/test_basic.py b/tests/test_basic.py
index ff860d4..bec6b54 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -151,9 +151,11 @@ def test_bytes_integer_signed_be():
 
 def test_enum():
     d = """
+    #define MyEnumVal 4
+
     enum Test16 : uint16 {
         A = 0x1,
-        B = 0x2
+        B = 0x6 - MyEnumVal
     };
 
     enum Test24 : uint24 {