Enum parsing

cstruct/__init__.py

@@ -41,8 +41,9 @@
 )
 from .abstract import CStructMeta, AbstractCStruct
 from .cstruct import CStruct
-from .c_parser import parse_def
+from .c_parser import parse_struct_def
 from .mem_cstruct import MemCStruct
+from .cenum import CEnum
 
 __all__ = [
     'LITTLE_ENDIAN',
@@ -150,7 +151,7 @@ def parse(
     """
     if __cls__ is None:
         __cls__ = CStruct
-    cls_def = parse_def(__struct__, __cls__=__cls__, **kargs)
+    cls_def = parse_struct_def(__struct__, __cls__=__cls__, **kargs)
     if cls_def is None:
         return None
     else:

cstruct/abstract.py

@@ -28,12 +28,15 @@
 from collections import OrderedDict
 from typing import Any, BinaryIO, List, Dict, Optional, Type, Tuple, Union
 import hashlib
-from .base import STRUCTS
-from .c_parser import parse_struct, parse_def, Tokens
+import sys
+from enum import IntEnum, EnumMeta, _EnumDict
+from unicodedata import name
+from .base import STRUCTS, ENUMS, DEFAULT_ENUM_SIZE
+from .c_parser import parse_struct, parse_struct_def, parse_enum_def,parse_enum, Tokens
 from .field import calculate_padding, FieldType
-from .exceptions import CStructException
+from .exceptions import CStructException, CEnumException
 
-__all__ = ['CStructMeta', 'AbstractCStruct']
+__all__ = ['CStructMeta', 'AbstractCStruct', "CEnumMeta", 'AbstractCEnum']
 
 
 class CStructMeta(ABCMeta):
@@ -48,7 +51,7 @@ def __new__(metacls: Type[type], name: str, bases: Tuple[str], namespace: Dict[s
                 namespace.update(parse_struct(**namespace))
             __struct__ = True
         if '__def__' in namespace:
-            namespace.update(parse_def(**namespace))
+            namespace.update(parse_struct_def(**namespace))
             __struct__ = True
         # Create the new class
         new_class: Type[Any] = super().__new__(metacls, name, bases, namespace)
@@ -130,7 +133,7 @@ def parse(
         cls_kargs['__struct__'] = __struct__
         if isinstance(__struct__, (str, Tokens)):
             del cls_kargs['__struct__']
-            cls_kargs.update(parse_def(__struct__, __cls__=cls, **cls_kargs))
+            cls_kargs.update(parse_struct_def(__struct__, __cls__=cls, **cls_kargs))
             cls_kargs['__struct__'] = None
         elif isinstance(__struct__, dict):
             del cls_kargs['__struct__']
@@ -257,3 +260,85 @@ def __setstate__(self, state: bytes) -> bool:
             state: bytes to be unpacked
         """
         return self.unpack(state)
+
+
+class CEnumMeta(EnumMeta):
+    class WrapperDict(_EnumDict):
+        def __setitem__(self, key: str, value: Any) -> None:
+            env = None
+            if key == "__enum__":
+                env = parse_enum(value)
+            elif key == "__def__":
+                env = parse_enum_def(value)
+
+            if env is not None:
+                # register the enum constants in the object namespace,
+                # using the Python Enum class Namespace dict that does the
+                # heavy lifting
+                for k, v in env["__constants__"].items():
+                    super().__setitem__(k, v)
+            else:
+                return super().__setitem__(key, value)
+
+    @classmethod
+    def __prepare__(metacls, cls, bases, **kwds):
+        namespace = EnumMeta.__prepare__(cls, bases, **kwds)
+        namespace.__class__ = metacls.WrapperDict
+        return namespace
+
+    def __new__(metacls: Type["CEnumMeta"], cls: str, bases: Tuple[Type, ...], classdict: _EnumDict, **kwds: Any) -> "CEnumMeta":
+        inst = super().__new__(metacls, cls, bases, classdict, **kwds)
+
+        if len(inst) > 0:
+            if "__size__" not in classdict:
+                inst.__size__ = DEFAULT_ENUM_SIZE
+                print(f"Warning: __size__ not specified for enum {cls}. Will default to {DEFAULT_ENUM_SIZE} bytes")
+            if not classdict.get("__anonymous__", False):
+                ENUMS[cls] = inst
+        return inst
+
+    @property
+    def size(cls) -> int:
+        "Enum size (in bytes)"
+        return cls.__size__
+
+class AbstractCEnum(IntEnum, metaclass=CEnumMeta):
+    """
+    Abstract C enum to Python class
+    """
+
+    @classmethod
+    def parse(
+        cls,
+        __enum__: Union[str, Tokens, Dict[str, Any]],
+        __name__: Optional[str] = None,
+        __size__: Optional[int] = None,
+        **kargs: Dict[str, Any]
+    ) -> Type["AbstractCEnum"]:
+        """
+        Return a new Python Enum class mapping a C enum definition
+
+        Args:
+            __enum__: Definition of the enum in C syntax
+            __name__: Name of the new Enum. If empty, a name based on the __enum__ hash is generated
+            __size__: Number of bytes that the enum should be read as
+
+        Returns:
+            cls: A new class mapping the definition
+        """
+
+        cls_kargs: Dict[str, Any] = dict(kargs)
+        if __size__ is not None:
+            cls_kargs['__size__'] = __size__
+
+        if isinstance(__enum__, (str, Tokens)):
+            cls_kargs.update(parse_enum_def(__enum__, __cls__=cls, **cls_kargs))
+        elif isinstance(__enum__, dict):
+            cls_kargs.update(__enum__)
+
+        if __name__ is None:
+            __name__ = cls.__name__ + "_" + hashlib.sha1(str(__enum__).encode("utf-8")).hexdigest()
+            cls_kargs["__anonymous__"] = True
+
+        cls_kargs.update(cls_kargs["__constants__"])
+        return cls(__name__, cls_kargs)

cstruct/base.py

@@ -25,7 +25,7 @@
 from typing import Any, Dict, Type, TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from .abstract import AbstractCStruct
+    from .abstract import AbstractCStruct, AbstractCEnum
 
 __all__ = [
     'LITTLE_ENDIAN',
@@ -48,6 +48,8 @@
 
 STRUCTS: Dict[str, Type["AbstractCStruct"]] = {}
 
+ENUMS: Dict[str, Type["AbstractCEnum"]] = {}
+
 DEFINES: Dict[str, Any] = {}
 
 TYPEDEFS: Dict[str, str] = {
@@ -91,4 +93,12 @@
     'uint64': 'Q',
 }
 
+ENUM_SIZE_TO_C_TYPE: Dict[int, str] = {
+    1: 'int8',
+    2: 'int16',
+    4: 'int32',
+    8: 'int64'
+}
+
 CHAR_ZERO = bytes('\0', 'ascii')
+DEFAULT_ENUM_SIZE = 4

cstruct/c_parser.py

@@ -25,15 +25,15 @@
 import re
 from collections import OrderedDict
 from typing import Union, Optional, Any, Dict, Type, TYPE_CHECKING
-from .base import DEFINES, TYPEDEFS, STRUCTS
+from .base import DEFINES, ENUMS, TYPEDEFS, STRUCTS
 from .field import calculate_padding, Kind, FieldType
 from .c_expr import c_eval
 from .exceptions import CStructException, ParserError
 
 if TYPE_CHECKING:
-    from .abstract import AbstractCStruct
+    from .abstract import AbstractCStruct, AbstractCEnum
 
-__all__ = ['parse_struct', 'parse_def', 'Tokens']
+__all__ = ['parse_struct', 'parse_struct_def', 'Tokens']
 
 
 class Tokens(object):
@@ -52,7 +52,7 @@ def __init__(self, text: str) -> None:
             else:
                 lines.append(line)
         text = " ".join(lines)
-        text = text.replace(";", " ; ").replace("{", " { ").replace("}", " } ")
+        text = text.replace(";", " ; ").replace("{", " { ").replace("}", " } ").replace(",", " , ").replace("=", " = ")
         self.tokens = text.split()
 
     def pop(self) -> str:
@@ -76,9 +76,9 @@ def parse_type(tokens: Tokens, __cls__: Type['AbstractCStruct'], byte_order: Opt
         raise ParserError("Parsing error")
     c_type = tokens.pop()
     # signed/unsigned/struct
-    if c_type in ['signed', 'unsigned', 'struct', 'union'] and len(tokens) > 1:
+    if c_type in ['signed', 'unsigned', 'struct', 'union', 'enum'] and len(tokens) > 1:
         c_type = c_type + " " + tokens.pop()
-    
+
     vlen = 1
     flexible_array = False
 
@@ -136,13 +136,32 @@ def parse_type(tokens: Tokens, __cls__: Type['AbstractCStruct'], byte_order: Opt
                 ref = STRUCTS[tail]
             except KeyError:
                 raise ParserError("Unknow {} {}".format(c_type, tail))
+    elif c_type.startswith('enum'):
+        #from .abstract import AbstractCEnum
+        from .cenum import CEnum
+
+        c_type, tail = c_type.split(' ', 1)
+        kind = Kind.ENUM
+        if tokens.get() == '{':  # Named nested struct
+            tokens.push(tail)
+            tokens.push(c_type)
+            ref = CEnum.parse(tokens, __name__=tail)
+        elif tail == '{':  # unnamed nested struct
+            tokens.push(tail)
+            tokens.push(c_type)
+            ref = CEnum.parse(tokens)
+        else:
+            try:
+                ref = ENUMS[tail]
+            except KeyError:
+                raise ParserError(f"Unknown '{c_type}' '{tail}'")
     else:  # other types
         kind = Kind.NATIVE
         ref = None
     return FieldType(kind, c_type, ref, vlen, flexible_array, byte_order, offset)
 
 
-def parse_def(
+def parse_struct_def(
     __def__: Union[str, Tokens],
     __cls__: Type['AbstractCStruct'],
     __byte_order__: Optional[str] = None,
@@ -169,6 +188,89 @@ def parse_def(
         raise ParserError("{} definition expected".format(vtype))
 
 
+def parse_enum_def(
+    __def__: Union[str, Tokens],
+    **kargs: Any
+) -> Optional[Dict[str, Any]]:
+    # naive C enum parsing
+    if isinstance(__def__, Tokens):
+        tokens = __def__
+    else:
+        tokens = Tokens(__def__)
+    if not tokens:
+        return None
+    kind = tokens.pop()
+    if kind not in ['enum']:
+        raise ParserError(f"enum expected - {kind}")
+
+    vtype = tokens.pop()
+    if tokens.get() == '{':  # named enum
+        tokens.pop()
+        return parse_enum(tokens)
+    elif vtype == '{':
+        return parse_enum(tokens)
+    else:
+        raise ParserError(f"{vtype} definition expected")
+
+
+def parse_enum(
+    __enum__: Union[str, Tokens],
+    **kargs: Any
+) -> Optional[Dict[str, Any]]:
+    constants: Dict[str, int] = OrderedDict()
+
+    if isinstance(__enum__, Tokens):
+        tokens = __enum__
+    else:
+        tokens = Tokens(__enum__)
+
+    while len(tokens):
+        if tokens.get() == '}':
+            tokens.pop()
+            break
+
+        name = tokens.pop()
+
+        next_token = tokens.pop()
+        if next_token in {",", "}"}:  # enum-constant without explicit value
+            if len(constants) == 0:
+                value = 0
+            else:
+                value = next(reversed(constants.values())) + 1
+        elif next_token == "=":  # enum-constant with explicit value
+            exp_elems = []
+            next_token = tokens.pop()
+            while next_token not in {",", "}"}:
+                exp_elems.append(next_token)
+                if len(tokens) > 0:
+                    next_token = tokens.pop()
+                else:
+                    break
+
+            if len(exp_elems) == 0:
+                raise ParserError(f"enum is missing value expression")
+
+            int_expr = " ".join(exp_elems)
+            try:
+                value = c_eval(int_expr)
+            except (ValueError, TypeError):
+                value = int(int_expr)
+        else:
+            raise ParserError(f"{__enum__} is not a valid enum expression")
+
+        if name in constants:
+            raise ParserError(f"duplicate enum name {name}")
+        constants[name] = value
+
+        if next_token == "}":
+            break
+
+    result = {
+        '__constants__': constants
+    }
+    return result
+
+
 def parse_struct(
     __struct__: Union[str, Tokens],
     __cls__: Type['AbstractCStruct'],

cstruct/cenum.py

@@ -0,0 +1,4 @@
+from .abstract import AbstractCEnum
+
+class CEnum(AbstractCEnum):
+  ...

cstruct/exceptions.py

@@ -23,11 +23,14 @@
 #
 
 __all__ = [
+    "CEnumException",
     "CStructException",
     "ParserError",
     "EvalError",
 ]
 
+class CEnumException(Exception):
+    pass
 
 class CStructException(Exception):
     pass