diff --git a/cstruct/__init__.py b/cstruct/__init__.py index d20062e..6dd3ec2 100644 --- a/cstruct/__init__.py +++ b/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: diff --git a/cstruct/abstract.py b/cstruct/abstract.py index f5aeca9..88e8493 100644 --- a/cstruct/abstract.py +++ b/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) \ No newline at end of file diff --git a/cstruct/base.py b/cstruct/base.py index b350a67..6b427ae 100644 --- a/cstruct/base.py +++ b/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 \ No newline at end of file diff --git a/cstruct/c_parser.py b/cstruct/c_parser.py index 476067a..3cad8f6 100644 --- a/cstruct/c_parser.py +++ b/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'], diff --git a/cstruct/cenum.py b/cstruct/cenum.py new file mode 100644 index 0000000..0fc1654 --- /dev/null +++ b/cstruct/cenum.py @@ -0,0 +1,4 @@ +from .abstract import AbstractCEnum + +class CEnum(AbstractCEnum): + ... \ No newline at end of file diff --git a/cstruct/exceptions.py b/cstruct/exceptions.py index 65cd668..3d66fb3 100644 --- a/cstruct/exceptions.py +++ b/cstruct/exceptions.py @@ -23,11 +23,14 @@ # __all__ = [ + "CEnumException", "CStructException", "ParserError", "EvalError", ] +class CEnumException(Exception): + pass class CStructException(Exception): pass diff --git a/cstruct/field.py b/cstruct/field.py index 2597b83..db4f6df 100644 --- a/cstruct/field.py +++ b/cstruct/field.py @@ -26,7 +26,7 @@ import struct from enum import Enum from typing import Optional, Any, List, Type, TYPE_CHECKING -from .base import NATIVE_ORDER, C_TYPE_TO_FORMAT +from .base import NATIVE_ORDER, C_TYPE_TO_FORMAT, ENUM_SIZE_TO_C_TYPE from .exceptions import ParserError if TYPE_CHECKING: @@ -58,6 +58,8 @@ class Kind(Enum): "Struct type" UNION = 2 "Union type" + ENUM = 3 + "Enum type" class FieldType(object): @@ -115,8 +117,12 @@ def unpack_from(self, buffer: bytes, offset: int = 0) -> Any: Returns: data: The unpacked data """ - if self.is_native: + if self.is_native or self.is_enum: result = struct.unpack_from(self.fmt, buffer, self.offset + offset) + + if self.is_enum: + result = tuple(map(self.ref, result)) + if self.is_array: return list(result) else: @@ -162,6 +168,11 @@ def is_native(self) -> bool: "True if the field is a native type (e.g. int, char)" return self.kind == Kind.NATIVE + @property + def is_enum(self) -> bool: + "True if the field is an enum" + return self.kind == Kind.ENUM + @property def is_struct(self) -> bool: "True if the field is a struct" @@ -180,13 +191,18 @@ def native_format(self) -> str: return C_TYPE_TO_FORMAT[self.c_type] except KeyError: raise ParserError("Unknow type {}".format(self.c_type)) + elif self.is_enum: + try: + return C_TYPE_TO_FORMAT[ENUM_SIZE_TO_C_TYPE[self.ref.size]] + except KeyError: + raise ParserError(f"Enum has invalid size. Needs to be in {ENUM_SIZE_TO_C_TYPE.keys()}") else: return 'c' @property def fmt(self) -> str: "Field format prefixed by byte order (struct library format)" - if self.is_native: + if self.is_native or self.is_enum: fmt = (str(self.vlen) if self.vlen > 1 or self.flexible_array else '') + self.native_format else: # Struct/Union fmt = str(self.vlen * self.ref.sizeof()) + self.native_format @@ -203,7 +219,7 @@ def vsize(self) -> int: @property def alignment(self) -> int: "Alignment" - if self.is_native: + if self.is_native or self.is_enum: if self.byte_order is not None: return struct.calcsize(self.byte_order + self.native_format) else: diff --git a/tests/test_cenum.py b/tests/test_cenum.py new file mode 100644 index 0000000..f137f8f --- /dev/null +++ b/tests/test_cenum.py @@ -0,0 +1,17 @@ +from cstruct import CEnum + +class Dummy(CEnum): + __enum__ = """ + A, + B, + C = 2, + D = 5 + 7, + E = 2 + """ + +def test_dummy(): + assert Dummy.A == 0 + assert Dummy.B == 1 + assert Dummy.C == 2 + assert Dummy.D == 12 + assert Dummy.E == 2 \ No newline at end of file