Add macros support

Signed-off-by: Cervenka Dusan <[email protected]>

Author: Hadatko

dissect/cstruct/cstruct.py

@@ -312,23 +312,23 @@ def values(self):
     def items(self):
         return self.typedefs.items()
 
-    def getEnumTypes(self):
-        retVal: list[EnumType] = []
+    def getTypes(self, classType):
+        retVal: list[classType] = []
         allTypes = self.typedefs
         for currentType in allTypes.values():
-            if isinstance(currentType, EnumType):
+            if isinstance(currentType, classType):
                 retVal.append(currentType)
 
         return retVal
 
+    def getEnumTypes(self):
+        return self.getTypes(EnumType)
+
     def getStructTypes(self):
-        retVal: list[StructureType] = []
-        allTypes = self.typedefs
-        for currentType in allTypes.values():
-            if isinstance(currentType, StructureType):
-                retVal.append(currentType)
+        return self.getTypes(StructureType)
 
-        return retVal
+    def getMacroTypes(self):
+        return self.getTypes(MacroType)
 
 
 class Parser(object):
@@ -360,10 +360,14 @@ class CStyleParser(Parser):
 
     def __init__(self, cstruct, compiled=True):
         self.compiled = compiled
+        self.clearVersionDetails()
         super(CStyleParser, self).__init__(cstruct)
 
-    # TODO: Implement proper parsing
+    def clearVersionDetails(self):
+        self.structVersionDetails = {"name":"", "valB":2**16-1, "valL":0}
+
     def parse(self, data):
+        self._macros(data)
         self._constants(data)
         self._enums(data)
         self._structs(data)
@@ -467,13 +471,68 @@ def _structs(self, data):
                     self.cstruct.addtype(td, st)
 
     def _parse_fields_struct(self, s):
+        def getNumber(sign:str, number:int):
+            retVal = number
+            direction = "="
+
+            if "<" in sign:
+                direction = "<"
+                if "=" not in sign:
+                    retVal -= 1
+            elif ">" in sign:
+                direction = ">"
+                if "=" not in sign:
+                    retVal += 1
+            elif "=" in sign:
+                pass
+            else:
+                retVal = None
+                direction = None
+
+            return retVal,direction
+
+        regex = r'(?P<if>#if\s*(?P<macroNameA>\w+)\s*(?P<signA>[<>=]*?)\s*(?P<valueA>[0-9]+)\s*(?:and\s*(?P<macroNameB>\w+)\s*(?P<signB>[<>=]*?)\s*(?P<valueB>[0-9]+))?)?\s*'+COMMENT_REGEX_START+r'(?P<type>[^\s]+)\s+(?P<name>[^\s\[:]+)(\s*:\s*(?P<bits>\d+))?(\[(?P<count>[^;\n]*)\])?;'+COMMENT_REGEX_END+r'(?P<endif>\s*#endif.*)?'
         fields = re.finditer(
-            COMMENT_REGEX_START+r'(?P<type>[^\s]+)\s+(?P<name>[^\s\[:]+)(\s*:\s*(?P<bits>\d+))?(\[(?P<count>[^;\n]*)\])?;'+COMMENT_REGEX_END,
+            regex,
             s, re.MULTILINE
         )
         r:list[StructureFieldType] = []
+
         for f in fields:
             d = f.groupdict()
+            # print(regex)
+            # print(d)
+            if d.get("if",None) != None:
+                macroNameA = d.get('macroNameA',None)
+                signA = d.get('signA',None)
+                valueA = d.get('valueA',None)
+                macroNameB = d.get('macroNameB', None)
+                signB = d.get('signB',None)
+                valueB = d.get('valueB', None)
+                if macroNameB != None and macroNameB != macroNameA:
+                    raise Exception(f"Expected macroNameA==macroNameB got {macroNameA}!={macroNameB}")
+                self.structVersionDetails["name"] = macroNameA
+                numA,dirA = getNumber(signA, valueA)
+                if dirA == "=":
+                    self.structVersionDetails["valL"] = valueA
+                    self.structVersionDetails["valB"] = valueA
+                else:
+                    if dirA == "<":
+                        self.structVersionDetails["valB"] = numA
+                    elif dirA == ">":
+                        self.structVersionDetails["valL"] = numA
+                    if signB != None:
+                        numB,dirB = getNumber(signB, valueB)
+                        if dirA == dirB:
+                            raise Exception(f"Expected dirA different sign <> than dirB got {dirA}=={dirB}")
+                        if dirB == "<":
+                            self.structVersionDetails["valB"] = numB
+                        elif dirB == ">":
+                            self.structVersionDetails["valL"] = numB
+                continue
+            if d.get("endif",None) != None:
+                self.clearVersionDetails()
+                continue
             if d['type'].startswith('//'):
                 continue
 
@@ -501,7 +560,7 @@ def _parse_fields_struct(self, s):
                 d['name'] = d['name'][1:]
                 type_ = Pointer(self.cstruct, type_)
 
-            field = StructureFieldType(d['name'], type_, int(d['bits']) if d['bits'] else None, commentAttributes=commentAttributes)
+            field = StructureFieldType(d['name'], type_, int(d['bits']) if d['bits'] else None, commentAttributes=commentAttributes, versionDetails=self.structVersionDetails.copy())
             r.append(field)
 
         return r
@@ -527,6 +586,19 @@ def parse_comment_block(self,s):
 
         return commentAttributes
 
+    def _macros(self, s):
+        r = re.finditer(
+            COMMENT_REGEX_START+r'#define\s*(?P<name>[A-Za-z0-9_]+)\s+(?P<value>[A-Za-z0-9_"()])'+COMMENT_REGEX_END,
+            s, re.MULTILINE
+        )
+
+        for t in r:
+            d = t.groupdict()
+
+            mt = MacroType(self.cstruct, d["name"], d["value"])
+
+            self.cstruct.addtype(d['name'], mt, True)
+
     def _lookups(self, data):
         r = re.finditer(r'\$(?P<name>[^\s]+) = ({[^}]+})\w*\n', data)
 
@@ -1081,12 +1153,13 @@ def reset(self):
 class StructureFieldType(object):
     """Holds a structure field."""
 
-    def __init__(self, name:str, type_:typeAll, bits=None, offset:typing.Union[int,None]=None, commentAttributes:typeCommentAttributes={}):
+    def __init__(self, name:str, type_:typeAll, bits=None, offset:typing.Union[int,None]=None, commentAttributes:typeCommentAttributes={}, versionDetails={}):
         self.name:str = name
         self.type:typeAll = type_
         self.bits = bits
         self.offset:typing.Union[int,None] = offset
         self.commentAttributes:typeCommentAttributes = commentAttributes
+        self.versionDetails = versionDetails
 
     def __repr__(self):
         return '<Field {} {} {}>'.format(self.name, self.type, self.commentAttributes)
@@ -1466,6 +1539,13 @@ def __repr__(self):
         return '<EnumType {}>'.format(self.name)
 
 
+class MacroType(RawType):
+    def __init__(self, cstruct, name:str, value:typing.Any):
+        self.name = name
+        self.value = value
+        super(MacroType, self).__init__(cstruct, name, 0)
+
+
 class EnumInstance(object):
     """Implements a value instance of an EnumType"""