This allow user to properly check datatype of Enums or Structures.

Signed-off-by: Cervenka Dusan <cervenka@acrios.com>

Author: Hadatko

dissect/cstruct/__init__.py

@@ -10,6 +10,9 @@
     Parser,
     RawType,
     BaseType,
+    StructureType,
+    EnumType,
+    ArrayType,
     Error,
     ParserError,
     CompilerError,
@@ -29,6 +32,9 @@
     "Parser",
     "RawType",
     "BaseType",
+    "StructureType",
+    "EnumType",
+    "ArrayType",
     "Error",
     "ParserError",
     "CompilerError",

dissect/cstruct/cstruct.py

@@ -66,7 +66,7 @@
 PRINTABLE = string.digits + string.ascii_letters + string.punctuation + " "
 
 COMPILE_TEMPL = """
-class {name}(Structure):
+class {name}(StructureType):
     def __init__(self, cstruct, structure, source=None):
         self.structure = structure
         self.source = source
@@ -85,7 +85,7 @@ def add_fields(self, name, type_, offset=None, commentAttributes={{}}):
         raise NotImplementedError("Can't add fields to a compiled structure")
 
     def __repr__(self):
-        return '<Structure {name} +compiled>'
+        return '<StructureType {name} +compiled>'
 """
 
 COMMENT_MULTILINE = r'/\*[^*]*\*+(?:[^/*][^*]*\*+)*/'
@@ -354,7 +354,7 @@ def _enums(self, data):
 
             commentAttributes = self.parse_comment_block(d['commentBlock'])
 
-            enum = Enum(
+            enum = EnumType(
                 self.cstruct, d['name'], self.cstruct.resolve(d['type']), values, valuesDetails, commentAttributes
             )
             self.cstruct.addtype(enum.name, enum)
@@ -407,7 +407,7 @@ def _structs(self, data):
             if d['type'] == 'struct':
                 data = self._parse_fields_struct(d['fields'][1:-1].strip())
                 commentAttributes = self.parse_comment_block(d['commentBlock'])
-                st = Structure(self.cstruct, name, data, commentAttributes)
+                st = StructureType(self.cstruct, name, data, commentAttributes)
                 if d['flags'] == 'compile' or self.compiled:
                     st = compiler.compile(st)
             elif d['typedef'] == 'typedef':
@@ -452,7 +452,7 @@ def _parse_fields_struct(self, s):
                     except Exception:
                         pass
 
-                type_ = Array(self.cstruct, type_, count)
+                type_ = ArrayType(self.cstruct, type_, count)
 
             if d['name'].startswith('*'):
                 d['name'] = d['name'][1:]
@@ -586,7 +586,7 @@ def _get(self):
         if self._value is None:
             pos = self._stream.tell()
             self._stream.seek(self._addr)
-            if isinstance(self._type, Array):
+            if isinstance(self._type, ArrayType):
                 r = self._type._read(self._stream, self._ctx)
             else:
                 r = self._type._read(self._stream)
@@ -761,7 +761,7 @@ def default_array(self):
         raise NotImplementedError()
 
     def __getitem__(self, count):
-        return Array(self.cstruct, self, count)
+        return ArrayType(self.cstruct, self, count)
 
     def __call__(self, *args, **kwargs):
         if len(args) > 0:
@@ -793,7 +793,7 @@ def __repr__(self):
         return BaseType.__repr__(self)
 
 
-class Structure(BaseType):
+class StructureType(BaseType):
     """Type class for structures."""
 
     def __init__(self, cstruct, name, fields=None, commentAttributes={}):
@@ -808,7 +808,7 @@ def __init__(self, cstruct, name, fields=None, commentAttributes={}):
             self.lookup[f.name] = f
 
         self._calc_offsets()
-        super(Structure, self).__init__(cstruct)
+        super(StructureType, self).__init__(cstruct)
 
     def _calc_offsets(self):
         offset = 0
@@ -860,7 +860,7 @@ def _calc_size(self):
         return size
 
     def _read(self, stream, *args, **kwargs):
-        log("[Structure::read] {} {}", self.name, self.size)
+        log("[StructureType::read] {} {}", self.name, self.size)
         bitbuffer = BitBuffer(stream, self.cstruct.endian)
 
         struct_start = stream.tell()
@@ -887,7 +887,7 @@ def _read(self, stream, *args, **kwargs):
             else:
                 bitbuffer.reset()
 
-            if isinstance(ft, (Array, Pointer)):
+            if isinstance(ft, (ArrayType, Pointer)):
                 v = ft._read(stream, r)
             else:
                 v = ft._read(stream)
@@ -950,7 +950,7 @@ def __len__(self):
         return self.size
 
     def __repr__(self):
-        return '<Structure {}>'.format(self.name)
+        return '<StructureType {}>'.format(self.name)
 
     def show(self, indent=0):
         """Pretty print this structure."""
@@ -965,7 +965,7 @@ def show(self, indent=0):
 
             print("{}+{} {} {}".format(' ' * indent, offset, field.name, field.type))
 
-            if isinstance(field.type, Structure):
+            if isinstance(field.type, StructureType):
                 field.type.show(indent + 1)
 
 
@@ -1037,7 +1037,7 @@ def __repr__(self):
         return '<Field {} {} {}>'.format(self.name, self.type, self.commentAttributes)
 
 
-class Array(BaseType):
+class ArrayType(BaseType):
     """Implements a fixed or dynamically sized array type.
 
     Example:
@@ -1053,7 +1053,7 @@ def __init__(self, cstruct, type_, count):
         self.count = count
         self.dynamic = isinstance(self.count, Expression) or self.count is None
 
-        super(Array, self).__init__(cstruct)
+        super(ArrayType, self).__init__(cstruct)
 
     def _read(self, stream, context=None):
         if self.count is None:
@@ -1332,10 +1332,10 @@ def default_array(self, count):
         return [0] * count
 
 
-class Enum(RawType):
-    """Implements an Enum type.
+class EnumType(RawType):
+    """Implements an EnumType type.
 
-    Enums can be made using any type. The API for accessing enums and their
+    EnumTypes can be made using any type. The API for accessing enums and their
     values is very similar to Python 3 native enums.
 
     Example:
@@ -1362,7 +1362,7 @@ def __init__(self, cstruct, name, type_, values, valuesDetails, commentAttribute
         for k, v in values.items():
             self.reverse[v] = k
 
-        super(Enum, self).__init__(cstruct, name, len(self.type))
+        super(EnumType, self).__init__(cstruct, name, len(self.type))
 
     def __call__(self, value):
         return EnumInstance(self, value)
@@ -1408,11 +1408,11 @@ def __contains__(self, attr):
         return attr in self.values
 
     def __repr__(self):
-        return '<Enum {}>'.format(self.name)
+        return '<EnumType {}>'.format(self.name)
 
 
 class EnumInstance(object):
-    """Implements a value instance of an Enum"""
+    """Implements a value instance of an EnumType"""
 
     def __init__(self, enum, value):
         self.enum = enum
@@ -1494,7 +1494,7 @@ def ctypes(structure):
         t = ctypes_type(field.type)
         fields.append((field.name, t))
 
-    tt = type(structure.name, (_ctypes.Structure, ), {"_fields_": fields})
+    tt = type(structure.name, (_ctypes.StructureType, ), {"_fields_": fields})
     return tt
 
 
@@ -1511,7 +1511,7 @@ def ctypes_type(t):
     if isinstance(t, CharType):
         return _ctypes.c_char
 
-    if isinstance(t, Array):
+    if isinstance(t, ArrayType):
         subtype = ctypes_type(t._type)
         return subtype * t.count
 
@@ -1534,7 +1534,7 @@ def compile(self, structure):
 
         env = {
             'OrderedDict': OrderedDict,
-            'Structure': Structure,
+            'StructureType': StructureType,
             'Instance': Instance,
             'Expression': Expression,
             'EnumInstance': EnumInstance,
@@ -1563,10 +1563,10 @@ def gen_struct_class(self, structure):
             if not isinstance(
                 ft,
                 (
-                    Structure,
+                    StructureType,
                     Pointer,
-                    Enum,
-                    Array,
+                    EnumType,
+                    ArrayType,
                     PackedType,
                     CharType,
                     WcharType,
@@ -1575,14 +1575,14 @@ def gen_struct_class(self, structure):
             ):
                 raise CompilerError("Unsupported type for compiler: {}".format(ft))
 
-            if isinstance(ft, Structure) or (
-                isinstance(ft, Array) and isinstance(ft.type, Structure)
+            if isinstance(ft, StructureType) or (
+                isinstance(ft, ArrayType) and isinstance(ft.type, StructureType)
             ):
                 if cur_block:
                     blocks.append(self.gen_read_block(read_size, cur_block))
 
                 struct_read = 's = stream.tell()\n'
-                if isinstance(ft, Array):
+                if isinstance(ft, ArrayType):
                     num = ft.count
 
                     if isinstance(num, Expression):
@@ -1686,17 +1686,17 @@ def gen_read_block(self, size, block):
             data_count = 1
             read_slice = ''
 
-            if isinstance(t, Enum):
+            if isinstance(t, EnumType):
                 t = t.type
             elif isinstance(t, Pointer):
                 t = self.cstruct.pointer
 
-            if isinstance(ft, Array):
+            if isinstance(ft, ArrayType):
                 count = t.count
                 data_count = count
                 t = t.type
 
-                if isinstance(t, Enum):
+                if isinstance(t, EnumType):
                     t = t.type
                 elif isinstance(t, Pointer):
                     t = self.cstruct.pointer
@@ -1720,7 +1720,7 @@ def gen_read_block(self, size, block):
                 else:
                     curtype = 'x'
 
-            if isinstance(t, (PackedType, CharType, WcharType, BytesInteger, Enum)):
+            if isinstance(t, (PackedType, CharType, WcharType, BytesInteger, EnumType)):
                 charcount = count
 
                 if isinstance(t, (CharType, WcharType, BytesInteger)):
@@ -1753,23 +1753,23 @@ def gen_read_block(self, size, block):
             else:
                 getter = 'data[{}]'.format(read_slice)
 
-            if isinstance(ft, Enum):
+            if isinstance(ft, EnumType):
                 getter = 'EnumInstance(self.cstruct.{type_name}, {getter})'.format(
                     type_name=ft.name, getter=getter
                 )
-            elif isinstance(ft, Array) and isinstance(ft.type, Enum):
+            elif isinstance(ft, ArrayType) and isinstance(ft.type, EnumType):
                 getter = '[EnumInstance(self.cstruct.{type_name}, d) for d in {getter}]'.format(
                     type_name=ft.type.name, getter=getter
                 )
             elif isinstance(ft, Pointer):
                 getter = 'PointerInstance(self.cstruct.{type_name}, stream, {getter}, r)'.format(
                     type_name=ft.type.name, getter=getter
                 )
-            elif isinstance(ft, Array) and isinstance(ft.type, Pointer):
+            elif isinstance(ft, ArrayType) and isinstance(ft.type, Pointer):
                 getter = '[PointerInstance(self.cstruct.{type_name}, stream, d, r) for d in {getter}]'.format(
                     type_name=ft.type.name, getter=getter
                 )
-            elif isinstance(ft, Array) and isinstance(t, PackedType):
+            elif isinstance(ft, ArrayType) and isinstance(t, PackedType):
                 getter = 'list({})'.format(getter)
 
             readcode.append(
@@ -1788,9 +1788,9 @@ def gen_read_block(self, size, block):
         return templ.format(''.join(fmt), '\n'.join(readcode))
 
     def gen_dynamic_block(self, field):
-        if not isinstance(field.type, Array):
+        if not isinstance(field.type, ArrayType):
             raise CompilerError(
-                "Only Array can be dynamic, got {!r}".format(field.type)
+                "Only ArrayType can be dynamic, got {!r}".format(field.type)
             )
 
         t = field.type.type
@@ -1928,8 +1928,8 @@ def dumpstruct(t, data=None, offset=0):
     Prints a colorized hexdump and parsed structure output.
 
     Args:
-        t: Structure or Instance to dump.
-        data: Bytes to parse the Structure on, if t is not a parsed Instance.
+        t: StructureType or Instance to dump.
+        data: Bytes to parse the StructureType on, if t is not a parsed Instance.
         offset: Byte offset of the hexdump.
     """
     colors = [
@@ -1946,7 +1946,7 @@ def dumpstruct(t, data=None, offset=0):
         g = t
         t = t._type
         data = g.dumps()
-    elif isinstance(t, Structure) and data:
+    elif isinstance(t, StructureType) and data:
         g = t(data)
     else:
         raise ValueError("Invalid arguments")