Set types as attributes on the cstruct object

Author: Schamper

dissect/cstruct/cstruct.py

@@ -56,9 +56,11 @@ def __init__(self, load: str = "", *, endian: str = "<", pointer: str | None = N
         self.endian = endian
 
         self.consts = {}
+        self.types = {}
         self.includes = []
+
         # fmt: off
-        self.typedefs = {
+        initial_types = {
             # Internal types
             "int8": self._make_packed_type("int8", "b", int),
             "uint8": self._make_packed_type("uint8", "B", int),
@@ -102,6 +104,22 @@ def __init__(self, load: str = "", *, endian: str = "<", pointer: str | None = N
             "signed long long": "int64",
             "unsigned long long": "uint64",
 
+            # Other convenience types
+            "u8": "uint8",
+            "u16": "uint16",
+            "u32": "uint32",
+            "u64": "uint64",
+            "u128": "uint128",
+            "__u8": "uint8",
+            "__u16": "uint16",
+            "__u32": "uint32",
+            "__u64": "uint64",
+            "__u128": "uint128",
+            "uchar": "uint8",
+            "ushort": "uint16",
+            "uint": "uint32",
+            "ulong": "uint32",
+
             # Windows types
             "BYTE": "uint8",
             "CHAR": "char",
@@ -169,25 +187,12 @@ def __init__(self, load: str = "", *, endian: str = "<", pointer: str | None = N
             "_DWORD": "uint32",
             "_QWORD": "uint64",
             "_OWORD": "uint128",
-
-            # Other convenience types
-            "u8": "uint8",
-            "u16": "uint16",
-            "u32": "uint32",
-            "u64": "uint64",
-            "u128": "uint128",
-            "__u8": "uint8",
-            "__u16": "uint16",
-            "__u32": "uint32",
-            "__u64": "uint64",
-            "__u128": "uint128",
-            "uchar": "uint8",
-            "ushort": "uint16",
-            "uint": "uint32",
-            "ulong": "uint32",
         }
         # fmt: on
 
+        for name, type_ in initial_types.items():
+            self.add_type(name, type_)
+
         pointer = pointer or ("uint64" if sys.maxsize > 2**32 else "uint32")
         self.pointer: type[BaseType] = self.resolve(pointer)
         self._anonymous_count = 0
@@ -202,27 +207,29 @@ def __getattr__(self, attr: str) -> Any:
             pass
 
         try:
-            return self.resolve(self.typedefs[attr])
+            return self.types[attr]
         except KeyError:
             pass
 
-        raise AttributeError(f"Invalid attribute: {attr}")
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute {attr!r}")
 
     def __copy__(self) -> cstruct:
         cs = cstruct(endian=self.endian, pointer=self.pointer.__name__)
         cs._anonymous_count = self._anonymous_count
-        cs.consts = self.consts.copy()
         cs.includes = self.includes.copy()
 
-        # Update typedefs to point to the new cstruct instance
-        for name, type in self.typedefs.items():
-            if name in cs.typedefs:
+        # Update types to point to the new cstruct instance
+        for name, type_ in self.types.items():
+            if name in cs.types:
                 continue
 
-            if isinstance(type, MetaType):
-                new_type = copy.copy(type)
+            if isinstance(type_, MetaType):
+                new_type = copy.copy(type_)
                 new_type.cs = cs
-                cs.typedefs[name] = new_type
+                cs.add_type(name, new_type)
+
+        for name, value in self.consts.items():
+            cs.add_const(name, value)
 
         return cs
 
@@ -231,22 +238,34 @@ def _next_anonymous(self) -> str:
         self._anonymous_count += 1
         return name
 
+    def _add_attr(self, name: str, value: Any) -> None:
+        # Names that collide with the cstruct class (e.g. a struct named ``load``) are not set as attributes
+        # to avoid breaking the instance. They remain accessible through ``resolve`` and the type dicts.
+        if name in _RESERVED_NAMES:
+            return
+        setattr(self, name, value)
+
     def add_type(self, name: str, type_: type[BaseType] | str, replace: bool = False) -> None:
-        """Add a type or type reference.
+        """Add a type or type alias.
 
         Only use this method when creating type aliases or adding already bound types.
+        All types will be resolved to their actual type objects prior to being added.
 
         Args:
             name: Name of the type to be added.
             type_: The type to be added. Can be a str reference to another type or a compatible type class.
+                   If a str is given, it will be resolved to the actual type object.
+            replace: Whether to replace the type if it already exists.
 
         Raises:
             ValueError: If the type already exists.
         """
-        if not replace and (name in self.typedefs and self.resolve(self.typedefs[name]) != self.resolve(type_)):
+        typeobj = self.resolve(type_)
+        if not replace and (existing := self.types.get(name)) is not None and existing is not typeobj:
             raise ValueError(f"Duplicate type: {name}")
 
-        self.typedefs[name] = type_
+        self.types[name] = typeobj
+        self._add_attr(name, typeobj)
 
     addtype = add_type
 
@@ -267,6 +286,28 @@ def add_custom_type(
         """
         self.add_type(name, self._make_type(name, (type_,), size, alignment=alignment, attrs=kwargs))
 
+    def add_const(self, name: str, value: Any) -> None:
+        """Add a constant value.
+
+        Args:
+            name: Name of the constant to be added.
+            value: The value of the constant.
+        """
+        self.consts[name] = value
+        self._add_attr(name, value)
+
+    def del_const(self, name: str) -> None:
+        """Delete a constant value.
+
+        Args:
+            name: Name of the constant to be deleted.
+
+        Raises:
+            KeyError: If the constant does not exist.
+        """
+        del self.consts[name]
+        self.__dict__.pop(name, None)
+
     def load(self, definition: str, deftype: int | None = None, **kwargs) -> cstruct:
         """Parse structures from the given definitions using the given definition type.
 
@@ -329,20 +370,13 @@ def resolve(self, name: type[BaseType] | str) -> type[BaseType]:
         Raises:
             ResolveError: If the type can't be resolved.
         """
-        type_name = name
-        if not isinstance(type_name, str):
-            return type_name
+        if not isinstance(name, str):
+            return name
 
-        for _ in range(10):
-            if type_name not in self.typedefs:
-                raise ResolveError(f"Unknown type {name}")
-
-            type_name = self.typedefs[type_name]
-
-            if not isinstance(type_name, str):
-                return type_name
-
-        raise ResolveError(f"Recursion limit exceeded while resolving type {name}")
+        try:
+            return self.types[name]
+        except KeyError:
+            raise ResolveError(f"Unknown type {name}") from None
 
     def copy(self) -> cstruct:
         """Create a copy of this cstruct instance.
@@ -519,6 +553,21 @@ class void(Void): ...
         # signed long long: TypeAlias = int64
         # unsigned long long: TypeAlias = uint64
 
+        u8: TypeAlias = uint8
+        u16: TypeAlias = uint16
+        u32: TypeAlias = uint32
+        u64: TypeAlias = uint64
+        u128: TypeAlias = uint128
+        __u8: TypeAlias = uint8
+        __u16: TypeAlias = uint16
+        __u32: TypeAlias = uint32
+        __u64: TypeAlias = uint64
+        __u128: TypeAlias = uint128
+        uchar: TypeAlias = uint8
+        ushort: TypeAlias = uint16
+        uint: TypeAlias = uint32
+        ulong: TypeAlias = uint32
+
         BYTE: TypeAlias = uint8
         CHAR: TypeAlias = char
         SHORT: TypeAlias = int16
@@ -584,20 +633,17 @@ class void(Void): ...
         _QWORD: TypeAlias = uint64
         _OWORD: TypeAlias = uint128
 
-        u8: TypeAlias = uint8
-        u16: TypeAlias = uint16
-        u32: TypeAlias = uint32
-        u64: TypeAlias = uint64
-        u128: TypeAlias = uint128
-        __u8: TypeAlias = uint8
-        __u16: TypeAlias = uint16
-        __u32: TypeAlias = uint32
-        __u64: TypeAlias = uint64
-        __u128: TypeAlias = uint128
-        uchar: TypeAlias = uint8
-        ushort: TypeAlias = uint16
-        uint: TypeAlias = uint32
-        ulong: TypeAlias = uint32
+
+# Attribute names that types and constants may never shadow: the public API and methods of the cstruct
+# class itself, plus the instance attributes created in __init__.
+_RESERVED_NAMES = frozenset(dir(cstruct)) | {
+    "endian",
+    "consts",
+    "types",
+    "includes",
+    "pointer",
+    "_anonymous_count",
+}
 
 
 def ctypes(structure: type[Structure]) -> type[_ctypes.Structure]:

dissect/cstruct/expression.py

@@ -12,6 +12,8 @@
     from dissect.cstruct import cstruct
     from dissect.cstruct.lexer import Token
 
+_MISSING = object()
+
 BINARY_OPERATORS: dict[TokenType, Callable[[int, int], int]] = {
     TokenType.PIPE: lambda a, b: a | b,
     TokenType.CARET: lambda a, b: a ^ b,
@@ -164,16 +166,13 @@ def evaluate(self, cs: cstruct, context: dict[str, int] | None = None) -> int:
                             obj = context[part]
                         elif part in cs.consts:
                             obj = cs.consts[part]
-                        elif part in cs.typedefs:
-                            obj = cs.resolve(part)
+                        elif part in cs.types:
+                            obj = cs.types[part]
                         else:
                             raise self._error(f"Unknown identifier: '{ident}'", token=token)
                     else:
-                        if isinstance(obj, dict) and part in obj:
-                            obj = obj[part]
-                        elif hasattr(obj, part):
-                            obj = getattr(obj, part)
-                        else:
+                        obj = obj.get(part, _MISSING) if isinstance(obj, dict) else getattr(obj, part, _MISSING)
+                        if obj is _MISSING:
                             raise self._error(f"Unknown identifier: '{ident}'", token=token)
 
                 try:

dissect/cstruct/parser.py

@@ -142,7 +142,8 @@ def _parse(self) -> None:
 
                 # If it's an anonymous enum/flag, add its members to the constants for convenience
                 if not type_.__name__:
-                    self.cs.consts.update(type_.__members__)
+                    for k, v in type_.__members__.items():
+                        self.cs.add_const(k, v)
 
                 self._expect(TokenType.SEMICOLON)
             else:
@@ -192,17 +193,17 @@ def _parse_define(self) -> None:
                             # If evaluation fails, just keep it as a string (e.g. for macro-like constants)
                             pass
 
-        self.cs.consts[name_token.value] = value
+        self.cs.add_const(name_token.value, value)
 
     def _parse_undef(self) -> None:
         """Parse an undef directive and remove the constant."""
         self._expect(TokenType.PP_UNDEF)
 
         name_token = self._expect(TokenType.IDENTIFIER)
-        if name_token.value in self.cs.consts:
-            del self.cs.consts[name_token.value]
-        else:
-            raise self._error(f"constant {name_token.value!r} not defined", token=name_token)
+        try:
+            self.cs.del_const(name_token.value)
+        except KeyError:
+            raise self._error(f"constant {name_token.value!r} not defined", token=name_token) from None
 
     def _parse_include(self) -> None:
         """Parse an include directive and add the included file to the includes list."""
@@ -527,7 +528,7 @@ def _parse_type_spec(self) -> type[BaseType]:
             ):
                 # This identifier is followed by a field delimiter, it should be the field name,
                 # UNLESS the current parts don't form a valid type yet.
-                if " ".join(parts) in self.cs.typedefs:
+                if " ".join(parts) in self.cs.types:
                     break
 
                 # Current parts don't resolve, consume and hope this completes the type name

dissect/cstruct/tools/stubgen.py

@@ -79,31 +79,29 @@ def generate_cstruct_stub(cs: cstruct, module_prefix: str = "", cls_name: str =
 
     defined_names = set()
 
-    # Then typedefs
-    for name, typedef in cs.typedefs.items():
-        if name in empty_cs.typedefs:
+    # Then types
+    for name, type_ in cs.types.items():
+        if name in empty_cs.types:
             continue
 
-        if typedef.__name__ in empty_cs.typedefs:
-            stub = f"{name}: TypeAlias = {cs_prefix}{typedef.__name__}"
-        elif typedef.__name__ in defined_names:
+        if type_.__name__ in empty_cs.types:
+            stub = f"{name}: TypeAlias = {cs_prefix}{type_.__name__}"
+        elif type_.__name__ in defined_names:
             # Create an alias to the type if we have already seen it before.
-            stub = f"{name}: TypeAlias = {typedef.__name__}"
-        elif issubclass(typedef, (types.Enum, types.Flag)):
-            stub = generate_enum_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
-        elif issubclass(typedef, types.Pointer):
-            typehint = generate_typehint(typedef, prefix=cs_prefix, module_prefix=module_prefix)
+            stub = f"{name}: TypeAlias = {type_.__name__}"
+        elif issubclass(type_, (types.Enum, types.Flag)):
+            stub = generate_enum_stub(type_, cs_prefix=cs_prefix, module_prefix=module_prefix)
+        elif issubclass(type_, types.Pointer):
+            typehint = generate_typehint(type_, prefix=cs_prefix, module_prefix=module_prefix)
             stub = f"{name}: TypeAlias = {typehint}"
-        elif issubclass(typedef, types.Structure):
-            stub = generate_structure_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
-        elif issubclass(typedef, types.BaseType):
-            stub = generate_generic_stub(typedef, cs_prefix=cs_prefix, module_prefix=module_prefix)
-        elif isinstance(typedef, str):
-            stub = f"{name}: TypeAlias = {typedef}"
+        elif issubclass(type_, types.Structure):
+            stub = generate_structure_stub(type_, cs_prefix=cs_prefix, module_prefix=module_prefix)
+        elif issubclass(type_, types.BaseType):
+            stub = generate_generic_stub(type_, cs_prefix=cs_prefix, module_prefix=module_prefix)
         else:
-            raise TypeError(f"Unknown typedef: {typedef}")
+            raise TypeError(f"Unknown type: {type_}")
 
-        defined_names.add(typedef.__name__)
+        defined_names.add(type_.__name__)
 
         body.append(textwrap.indent(stub, prefix=indent))
 
@@ -158,7 +156,6 @@ def generate_structure_stub(
     module_prefix: str = "",
 ) -> str:
     result = [f"class {name_prefix}{structure.__name__}({module_prefix}{structure.__base__.__name__}):"]
-
     indent = " " * 4
 
     args = ["self"]
@@ -171,7 +168,7 @@ def generate_structure_stub(
         while issubclass(nested_type, types.BaseArray):
             nested_type = nested_type.type
 
-        if issubclass(nested_type, types.Structure) and type_name not in structure.cs.typedefs:
+        if issubclass(nested_type, types.Structure) and type_name not in structure.cs.types:
             inlined = True
             inline_stub = generate_structure_stub(nested_type, cs_prefix=cs_prefix, module_prefix=module_prefix)