Merge pull request #7 from ACRIOS-Systems/feature/necessaryTyping

Feature/necessary typing

Author: Hadatko

dissect/cstruct/__init__.py

@@ -11,8 +11,15 @@
     RawType,
     BaseType,
     StructureType,
+    StructureFieldType,
     EnumType,
     ArrayType,
+    PackedType,
+    CharType,
+    WcharType,
+    UnionType,
+    VoidType,
+    typeAll,
     Error,
     ParserError,
     CompilerError,
@@ -33,8 +40,15 @@
     "RawType",
     "BaseType",
     "StructureType",
+    "StructureFieldType",
     "EnumType",
     "ArrayType",
+    "PackedType",
+    "CharType",
+    "WcharType",
+    "UnionType",
+    "VoidType",
+    "typeAll",
     "Error",
     "ParserError",
     "CompilerError",

dissect/cstruct/cstruct.py

@@ -33,6 +33,7 @@
 from io import BytesIO
 from collections import OrderedDict
 import json
+import typing
 
 try:
     from builtins import bytes as newbytes
@@ -88,6 +89,7 @@ def __repr__(self):
         return '<StructureType {name} +compiled>'
 """
 
+#common regex values
 COMMENT_MULTILINE = r'/\*[^*]*\*+(?:[^/*][^*]*\*+)*/'
 COMMENT_INLINE = r'//[^\n]*'
 COMMENT_MULTILINE_REPEATED = r'(^[ ]*('+COMMENT_INLINE+r'|'+COMMENT_MULTILINE+r'([ ]*('+COMMENT_INLINE+r'|'+COMMENT_MULTILINE+r'))*)[ \t\r]*\n?)*^[ ]*('+COMMENT_INLINE+r'|(?P<commentBlock>'+COMMENT_MULTILINE+r'))+'
@@ -97,6 +99,12 @@ def __repr__(self):
 #print(f"COMMENT_REGEX_START:{COMMENT_REGEX_START}")
 #print(f"COMMENT_REGEX_END:{COMMENT_REGEX_END}")
 
+#types
+typeCommentAttributes=dict[str,str]
+typeValuesDetails =dict[str,dict[str,typing.Union[str,int,typeCommentAttributes]]]
+typeAll = typing.Union[str,'PackedType','CharType','WcharType','BytesInteger','VoidType','ArrayType','StructureType','EnumType']
+typeTypedefs = dict[str, typeAll]
+
 class Error(Exception):
     pass
 
@@ -139,7 +147,7 @@ def __init__(self, endian='<', pointer='uint64'):
 
         self.consts = {}
         self.lookups = {}
-        self.typedefs = {
+        self.typedefs:typeTypedefs = {
             'byte': 'int8',
             'ubyte': 'uint8',
             'uchar': 'uint8',
@@ -186,7 +194,7 @@ def __init__(self, endian='<', pointer='uint64'):
 
         self.pointer = self.resolve(pointer)
 
-    def addtype(self, name, t, replace=False):
+    def addtype(self, name:str, t:typeAll, replace=False):
         """Add a type or type reference.
 
         Args:
@@ -197,8 +205,7 @@ def addtype(self, name, t, replace=False):
         Raises:
             ValueError: If the type already exists.
         """
-        name = name.lower()
-        if not replace and name.lower() in self.typedefs:
+        if not replace and name in self.typedefs:
             raise ValueError("Duplicate type: %s" % name)
 
         self.typedefs[name] = t
@@ -223,7 +230,7 @@ def load(self, s, deftype=None, **kwargs):
 
         parser.parse(s)
 
-    def loadfile(self, s, deftype=None, **kwargs):
+    def loadfile(self, s:str, deftype:typing.Optional[int]=None, **kwargs:typing.Union['cstruct',bool]):
         """Load structure definitions from a file.
 
         The given path will be read and parsed using the .load() function.
@@ -268,25 +275,61 @@ def resolve(self, name):
             return t
 
         for i in xrange(10):
-            if t.lower() not in self.typedefs:
+            if t not in self.typedefs:
                 raise ResolveError("Unknown type %s" % name)
 
-            t = self.typedefs[t.lower()]
+            t = self[t]
 
             if not isinstance(t, str):
                 return t
 
         raise ResolveError("Recursion limit exceeded while resolving type %s" % name)
 
     def __getattr__(self, attr):
-        if attr.lower() in self.typedefs:
-            return self.typedefs[attr.lower()]
+        if attr in self.typedefs:
+            return self.typedefs[attr]
 
         if attr in self.consts:
             return self.consts[attr]
 
+        for enum in self.getEnumTypes():
+            for enumMember in enum.values.items():
+                if enumMember[0] == attr:
+                    return enumMember[1]
+
+
         raise AttributeError("Invalid Attribute: %s" % attr)
 
+    def __getitem__(self, item):
+        return self.__getattr__(item)
+
+    def keys(self):
+        return self.typedefs.keys()
+
+    def values(self):
+        return self.typedefs.values()
+
+    def items(self):
+        return self.typedefs.items()
+
+    def getEnumTypes(self):
+        retVal: list[EnumType] = []
+        allTypes = self.typedefs
+        for currentType in allTypes.values():
+            if isinstance(currentType, EnumType):
+                retVal.append(currentType)
+
+        return retVal
+
+    def getStructTypes(self):
+        retVal: list[StructureType] = []
+        allTypes = self.typedefs
+        for currentType in allTypes.values():
+            if isinstance(currentType, StructureType):
+                retVal.append(currentType)
+
+        return retVal
+
 
 class Parser(object):
     """Base class for definition parsers.
@@ -428,7 +471,7 @@ def _parse_fields_struct(self, s):
             COMMENT_REGEX_START+r'(?P<type>[^\s]+)\s+(?P<name>[^\s\[:]+)(\s*:\s*(?P<bits>\d+))?(\[(?P<count>[^;\n]*)\])?;'+COMMENT_REGEX_END,
             s, re.MULTILINE
         )
-        r = []
+        r:list[StructureFieldType] = []
         for f in fields:
             d = f.groupdict()
             if d['type'].startswith('//'):
@@ -458,13 +501,13 @@ def _parse_fields_struct(self, s):
                 d['name'] = d['name'][1:]
                 type_ = Pointer(self.cstruct, type_)
 
-            field = StructureField(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)
             r.append(field)
 
         return r
 
     def parse_comment_block(self,s):
-        commentAttributes={}
+        commentAttributes:typeCommentAttributes={}
 
         #parse the comment header
         if s is not None and s.startswith('/*'):
@@ -622,7 +665,7 @@ class Expression(object):
         ('<<', lambda a, b: a << b),
     ]
 
-    def __init__(self, cstruct, expr):
+    def __init__(self, cstruct:cstruct, expr):
         self.cstruct = cstruct
         self.expr = expr
 
@@ -667,6 +710,11 @@ def evaluate_part(self, e, v):
         if e in self.cstruct.consts:
             return self.cstruct.consts[e]
 
+        for enum in self.cstruct.getEnumTypes():
+            for enumMember in enum.values.items():
+                if enumMember[0] == e:
+                    return enumMember[1]
+
         return int(e)
 
     def __repr__(self):
@@ -778,9 +826,9 @@ def __call__(self, *args, **kwargs):
 class RawType(BaseType):
     """Base class for raw types that have a name and size."""
 
-    def __init__(self, cstruct, name=None, size=0):
-        self.name = name
-        self.size = size
+    def __init__(self, cstruct, name:typing.Union[str,None]=None, size:int=0):
+        self.name:typing.Union[str,None] = name
+        self.size:int = size
         super(RawType, self).__init__(cstruct)
 
     def __len__(self):
@@ -796,12 +844,12 @@ def __repr__(self):
 class StructureType(BaseType):
     """Type class for structures."""
 
-    def __init__(self, cstruct, name, fields=None, commentAttributes={}):
-        self.name = name
-        self.size = None
+    def __init__(self, cstruct, name:str, fields:list['StructureFieldType']=[], commentAttributes:typeCommentAttributes={}):
+        self.name:str = name
+        self.size:typing.Union[int,None] = None
         self.lookup = OrderedDict()
-        self.fields = fields if fields else []
-        self.commentAttributes = commentAttributes
+        self.fields:list[StructureFieldType] = fields
+        self.commentAttributes:typeCommentAttributes = commentAttributes
 
 
         for f in self.fields:
@@ -917,15 +965,15 @@ def _write(self, stream, data):
 
         return num
 
-    def add_fields(self, name, type_, offset=None, commentAttributes={}):
+    def add_fields(self, name, type_, offset=None, commentAttributes:typeCommentAttributes={}):
         """Add a field to this structure.
 
         Args:
             name: The field name.
             type_: The field type.
             offset: The field offset.
         """
-        field = StructureField(name, type_, offset=offset, commentAttributes=commentAttributes)
+        field = StructureFieldType(name, type_, offset=offset, commentAttributes=commentAttributes)
         self.fields.append(field)
         self.lookup[name] = field
         self.size = None
@@ -1023,15 +1071,15 @@ def reset(self):
         self._remaining = 0
 
 
-class StructureField(object):
+class StructureFieldType(object):
     """Holds a structure field."""
 
-    def __init__(self, name, type_, bits=None, offset=None, commentAttributes={}):
-        self.name = name
-        self.type = type_
+    def __init__(self, name:str, type_:typeAll, bits=None, offset:typing.Union[int,None]=None, commentAttributes:typeCommentAttributes={}):
+        self.name:str = name
+        self.type:typeAll = type_
         self.bits = bits
-        self.offset = offset
-        self.commentAttributes = commentAttributes
+        self.offset:typing.Union[int,None] = offset
+        self.commentAttributes:typeCommentAttributes = commentAttributes
 
     def __repr__(self):
         return '<Field {} {} {}>'.format(self.name, self.type, self.commentAttributes)
@@ -1048,8 +1096,8 @@ class ArrayType(BaseType):
             x[expr] -> expr -> dynamic length.
     """
 
-    def __init__(self, cstruct, type_, count):
-        self.type = type_
+    def __init__(self, cstruct, type_:typeAll, count):
+        self.type:typeAll = type_
         self.count = count
         self.dynamic = isinstance(self.count, Expression) or self.count is None
 
@@ -1352,11 +1400,11 @@ class EnumType(RawType):
             };
     """
 
-    def __init__(self, cstruct, name, type_, values, valuesDetails, commentAttributes={}):
-        self.type = type_
-        self.values = values
-        self.valuesDetails = valuesDetails
-        self.commentAttributes = commentAttributes
+    def __init__(self, cstruct, name:str, type_:typeAll, values:dict[str,int], valuesDetails:typeValuesDetails, commentAttributes:typeCommentAttributes={}):
+        self.type:typeAll = type_
+        self.values:dict[str,int] = values
+        self.valuesDetails:typeValuesDetails = valuesDetails
+        self.commentAttributes:typeCommentAttributes = commentAttributes
         self.reverse = {}
 
         for k, v in values.items():
@@ -1449,10 +1497,10 @@ def __call__(self):
         return self.value
 
 
-class Union(RawType):
+class UnionType(RawType):
     def __init__(self, cstruct):
         self.cstruct = cstruct
-        super(Union, self).__init__(cstruct)
+        super(UnionType, self).__init__(cstruct)
 
     def _read(self, stream):
         raise NotImplementedError()