Set types as attributes on the cstruct object

Author: Schamper

dissect/cstruct/cstruct.py

@@ -48,8 +48,11 @@ def __init__(self, endian: str = "<", pointer: str | None = None):
 
         self.consts = {}
         self.lookups = {}
+        self.types = {}
+        self.typedefs = {}
         # fmt: off
-        self.typedefs = {
+
+        initial_types = {
             # Internal types
             "int8": self._make_packed_type("int8", "b", int),
             "uint8": self._make_packed_type("uint8", "B", int),
@@ -93,6 +96,21 @@ def __init__(self, endian: str = "<", pointer: str | None = None):
             "signed long long": "long long",
             "unsigned long long": "uint64",
 
+            # Other convenience types
+            "u1": "uint8",
+            "u2": "uint16",
+            "u4": "uint32",
+            "u8": "uint64",
+            "u16": "uint128",
+            "__u8": "uint8",
+            "__u16": "uint16",
+            "__u32": "uint32",
+            "__u64": "uint64",
+            "uchar": "uint8",
+            "ushort": "unsigned short",
+            "uint": "unsigned int",
+            "ulong": "unsigned long",
+
             # Windows types
             "BYTE": "uint8",
             "CHAR": "char",
@@ -160,24 +178,12 @@ def __init__(self, endian: str = "<", pointer: str | None = None):
             "_DWORD": "uint32",
             "_QWORD": "uint64",
             "_OWORD": "uint128",
-
-            # Other convenience types
-            "u1": "uint8",
-            "u2": "uint16",
-            "u4": "uint32",
-            "u8": "uint64",
-            "u16": "uint128",
-            "__u8": "uint8",
-            "__u16": "uint16",
-            "__u32": "uint32",
-            "__u64": "uint64",
-            "uchar": "uint8",
-            "ushort": "unsigned short",
-            "uint": "unsigned int",
-            "ulong": "unsigned long",
         }
         # 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 = self.resolve(pointer)
         self._anonymous_count = 0
@@ -188,37 +194,71 @@ def __getattr__(self, attr: str) -> Any:
         except KeyError:
             pass
 
+        try:
+            return self.types[attr]
+        except KeyError:
+            pass
+
         try:
             return self.resolve(self.typedefs[attr])
         except KeyError:
             pass
 
-        raise AttributeError(f"Invalid attribute: {attr}")
+        return super().__getattribute__(attr)
 
     def _next_anonymous(self) -> str:
         name = f"__anonymous_{self._anonymous_count}__"
         self._anonymous_count += 1
         return name
 
+    def _add_attr(self, name: str, value: Any, replace: bool = False) -> None:
+        if not replace and (name in self.__dict__ and self.__dict__[name] != value):
+            raise ValueError(f"Attribute already exists: {name}")
+        setattr(self, name, value)
+
     def add_type(self, name: str, type_: MetaType | str, replace: bool = False) -> None:
         """Add a type or type reference.
 
         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.
+        Use :func:`add_typedef` to add type references.
 
         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.
 
         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 (name in self.types and self.types[name] != typeobj):
             raise ValueError(f"Duplicate type: {name}")
 
-        self.typedefs[name] = type_
+        self.types[name] = typeobj
+        self._add_attr(name, typeobj, replace=replace)
 
     addtype = add_type
 
+    def add_typedef(self, name: str, type_: str, replace: bool = False) -> None:
+        """Add a type reference.
+
+        Use this method to add type references to this cstruct instance. This is largely a convenience method for the
+        internal :func:`add_type` method.
+
+        Args:
+            name: Name of the type to be added.
+            type_: The type reference to be added.
+            replace: Whether to replace the type if it already exists.
+        """
+        if not isinstance(type_, str):
+            raise TypeError("Type reference must be a string")
+
+        if not replace and (name in self.typedefs and self.resolve(self.typedefs[name]) != self.resolve(type_)):
+            raise ValueError(f"Duplicate type: {name}")
+
+        self.typedefs[name] = type_
+
     def add_custom_type(
         self, name: str, type_: MetaType, size: int | None = None, alignment: int | None = None, **kwargs
     ) -> None:
@@ -236,6 +276,16 @@ 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, replace=True)
+
     def load(self, definition: str, deftype: int | None = None, **kwargs) -> cstruct:
         """Parse structures from the given definitions using the given definition type.
 
@@ -307,14 +357,14 @@ def resolve(self, name: str) -> MetaType:
             return type_name
 
         for _ in range(10):
+            if type_name in self.types:
+                return self.types[type_name]
+
             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}")
 
     def _make_type(

dissect/cstruct/parser.py

@@ -15,6 +15,7 @@
 
 if TYPE_CHECKING:
     from dissect.cstruct import cstruct
+    from dissect.cstruct.types.structure import Structure
 
 
 class Parser:
@@ -96,7 +97,7 @@ def _constant(self, tokens: TokenConsumer) -> None:
             except (ExpressionParserError, ExpressionTokenizerError):
                 pass
 
-        self.cstruct.consts[match["name"]] = value
+        self.cstruct.add_const(match["name"], value)
 
     def _enum(self, tokens: TokenConsumer) -> None:
         # We cheat with enums because the entire enum is in the token
@@ -137,18 +138,21 @@ def _enum(self, tokens: TokenConsumer) -> None:
 
         enum = factory(d["name"] or "", self.cstruct.resolve(d["type"]), values)
         if not enum.__name__:
-            self.cstruct.consts.update(enum.__members__)
+            for k, v in enum.__members__.items():
+                self.cstruct.add_const(k, v)
         else:
             self.cstruct.add_type(enum.__name__, enum)
 
         tokens.eol()
 
     def _typedef(self, tokens: TokenConsumer) -> None:
         tokens.consume()
+        type_name = None
         type_ = None
 
         if tokens.next == self.TOK.IDENTIFIER:
-            type_ = self.cstruct.resolve(self._identifier(tokens))
+            type_name = self._identifier(tokens)
+            type_ = self.cstruct.resolve(type_name)
         elif tokens.next == self.TOK.STRUCT:
             # The register thing is a bit dirty
             # Basically consumes all NAME tokens and
@@ -157,12 +161,15 @@ def _typedef(self, tokens: TokenConsumer) -> None:
 
         names = self._names(tokens)
         for name in names:
-            type_, name, bits = self._parse_field_type(type_, name)
+            new_type, name, bits = self._parse_field_type(type_, name)
             if bits is not None:
                 raise ParserError(f"line {self._lineno(tokens.previous)}: typedefs cannot have bitfields")
-            self.cstruct.add_type(name, type_)
+            if type_name is None or new_type is not type_:
+                self.cstruct.add_type(name, new_type)
+            else:
+                self.cstruct.add_typedef(name, type_name)
 
-    def _struct(self, tokens: TokenConsumer, register: bool = False) -> None:
+    def _struct(self, tokens: TokenConsumer, register: bool = False) -> type[Structure]:
         stype = tokens.consume()
 
         factory = self.cstruct._make_union if stype.value.startswith("union") else self.cstruct._make_struct
@@ -399,7 +406,7 @@ def _constants(self, data: str) -> None:
             except (ValueError, SyntaxError):
                 pass
 
-            self.cstruct.consts[d["name"]] = v
+            self.cstruct.add_const(d["name"], v)
 
     def _enums(self, data: str) -> None:
         r = re.finditer(
@@ -481,7 +488,7 @@ def _structs(self, data: str) -> None:
             if d["defs"]:
                 for td in d["defs"].strip().split(","):
                     td = td.strip()
-                    self.cstruct.add_type(td, st)
+                    self.cstruct.add_typedef(td, st)
 
     def _parse_fields(self, data: str) -> None:
         fields = re.finditer(

dissect/cstruct/types/structure.py

@@ -241,7 +241,7 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> Struc
         sizes = {}
         for field in cls.__fields__:
             offset = stream.tell()
-            field_type = cls.cs.resolve(field.type)
+            field_type = field.type
 
             if field.offset is not None and offset != struct_start + field.offset:
                 # Field is at a specific offset, either alligned or added that way
@@ -295,7 +295,7 @@ def _write(cls, stream: BinaryIO, data: Structure) -> int:
         num = 0
 
         for field in cls.__fields__:
-            field_type = cls.cs.resolve(field.type)
+            field_type = field.type
 
             bit_field_type = (
                 (field_type.type if isinstance(field_type, EnumMetaType) else field_type) if field.bits else None
@@ -460,7 +460,7 @@ def _read_fields(
             buf = io.BytesIO(stream.read(cls.size))
 
         for field in cls.__fields__:
-            field_type = cls.cs.resolve(field.type)
+            field_type = field.type
 
             start = 0
             if field.offset is not None:

tests/test_basic.py

@@ -31,7 +31,7 @@ def test_load_file(cs: cstruct, compiled: bool) -> None:
     path = Path(__file__).parent / "data/testdef.txt"
 
     cs.loadfile(path, compiled=compiled)
-    assert "test" in cs.typedefs
+    assert "test" in cs.types
 
 
 def test_read_type_name(cs: cstruct) -> None:
@@ -46,7 +46,7 @@ def test_type_resolve(cs: cstruct) -> None:
 
     cs.add_type("ref0", "uint32")
     for i in range(1, 15):  # Recursion limit is currently 10
-        cs.add_type(f"ref{i}", f"ref{i - 1}")
+        cs.add_typedef(f"ref{i}", f"ref{i - 1}")
 
     with pytest.raises(ResolveError, match="Recursion limit exceeded"):
         cs.resolve("ref14")
@@ -404,7 +404,7 @@ def test_reserved_keyword(cs: cstruct, compiled: bool) -> None:
     cs.load(cdef, compiled=compiled)
 
     for name in ["in", "class", "for"]:
-        assert name in cs.typedefs
+        assert name in cs.types
         assert verify_compiled(cs.resolve(name), compiled)
 
         assert cs.resolve(name)(b"\x01").a == 1