Add examples and unit tests for custom types in cstruct v4 (#83)

Author: Miauwkeru

dissect/cstruct/types/base.py

@@ -5,10 +5,10 @@
 from typing import TYPE_CHECKING, Any, BinaryIO, Callable
 
 from dissect.cstruct.exceptions import ArraySizeError
+from dissect.cstruct.expression import Expression
 
 if TYPE_CHECKING:
     from dissect.cstruct.cstruct import cstruct
-    from dissect.cstruct.expression import Expression
 
 
 EOF = -0xE0F  # Negative counts are illegal anyway, so abuse that for our EOF sentinel
@@ -210,6 +210,13 @@ class BaseType(metaclass=MetaType):
     dumps = _overload(MetaType.dumps)
     write = _overload(MetaType.write)
 
+    def __len__(self) -> int:
+        """Return the byte size of the type."""
+        if self.__class__.size is None:
+            raise TypeError("Dynamic size")
+
+        return self.__class__.size
+
 
 class ArrayMetaType(MetaType):
     """Base metaclass for array-like types."""
@@ -222,15 +229,17 @@ def _read(cls, stream: BinaryIO, context: dict[str, Any] = None) -> Array:
         if cls.null_terminated:
             return cls.type._read_0(stream, context)
 
-        if cls.dynamic:
+        if isinstance(cls.num_entries, int):
+            num = max(0, cls.num_entries)
+        elif cls.num_entries is None:
+            num = EOF
+        elif isinstance(cls.num_entries, Expression):
             try:
                 num = max(0, cls.num_entries.evaluate(context))
             except Exception:
                 if cls.num_entries.expression != "EOF":
                     raise
                 num = EOF
-        else:
-            num = max(0, cls.num_entries)
 
         return cls.type._read_array(stream, num, context)
 

examples/protobuf.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+from typing import Any, BinaryIO
+
+from dissect.cstruct import cstruct
+from dissect.cstruct.types import BaseType
+
+
+class ProtobufVarint(BaseType):
+    """Implements a protobuf integer type for dissect.cstruct that can span a variable amount of bytes.
+
+    Mainly follows the BaseType implementation with minor tweaks
+    to support protobuf's msb varint implementation.
+
+    Resources:
+        - https://protobuf.dev/programming-guides/encoding/
+        - https://github.com/protocolbuffers/protobuf/blob/main/python/google/protobuf/internal/decoder.py
+    """
+
+    @classmethod
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> int:
+        return decode_varint(stream)
+
+    @classmethod
+    def _write(cls, stream: BinaryIO, data: int) -> int:
+        return stream.write(encode_varint(data))
+
+
+def decode_varint(stream: BinaryIO) -> int:
+    """Reads a varint from the provided buffer stream.
+
+    If we have not reached the end of a varint, the msb will be 1.
+    We read every byte from our current position until the msb is 0.
+    """
+    result = 0
+    i = 0
+    while True:
+        byte = stream.read(1)
+        result |= (byte[0] & 0x7F) << (i * 7)
+        i += 1
+        if byte[0] & 0x80 == 0:
+            break
+
+    return result
+
+
+def encode_varint(number: int) -> bytes:
+    """Encode a decoded protobuf varint to its original bytes."""
+    buf = []
+    while True:
+        towrite = number & 0x7F
+        number >>= 7
+        if number:
+            buf.append(towrite | 0x80)
+        else:
+            buf.append(towrite)
+            break
+    return bytes(buf)
+
+
+if __name__ == "__main__":
+    cdef = """
+    struct foo {
+        uint32 foo;
+        varint size;
+        char   bar[size];
+    };
+    """
+
+    cs = cstruct(endian=">")
+    cs.add_custom_type("varint", ProtobufVarint)
+    cs.load(cdef, compiled=False)
+
+    aaa = b"a" * 123456
+    buf = b"\x00\x00\x00\x01\xc0\xc4\x07" + aaa
+    foo = cs.foo(buf + b"\x01\x02\x03")
+    assert foo.foo == 1
+    assert foo.size == 123456
+    assert foo.bar == aaa
+    assert foo.dumps() == buf
+
+    assert cs.varint[2](b"\x80\x01\x80\x02") == [128, 256]
+    assert cs.varint[2].dumps([128, 256]) == b"\x80\x01\x80\x02"

examples/secdesc.py

@@ -52,8 +52,8 @@
 struct ACCESS_ALLOWED_OBJECT_ACE {
     uint32  Mask;
     uint32  Flags;
-    char    ObjectType[Flags & 1 * 16];
-    char    InheritedObjectType[Flags & 2 * 8];
+    char    ObjectType[(Flags & 1) * 16];
+    char    InheritedObjectType[(Flags & 2) * 8];
     LDAP_SID Sid;
 };
 """
@@ -97,12 +97,9 @@ def __init__(self, fh=None, in_obj=None):
             self.ldap_sid = in_obj
 
     def __repr__(self):
-        return "S-{}-{}-{}".format(
-            self.ldap_sid.Revision,
-            bytearray(self.ldap_sid.IdentifierAuthority.Value)[5],
-            "-".join(["{}".format(v) for v in self.ldap_sid.SubAuthority]),
-        )
-
+        authority = bytearray(self.ldap_sid.IdentifierAuthority.Value)[5]
+        sub_authority = "-".join(f"{v}" for v in self.ldap_sid.SubAuthority)
+        return f"S-{self.ldap_sid.Revision}-{authority}-{sub_authority}"
 
 class ACL:
     def __init__(self, fh):

tests/test_types_custom.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import Any, BinaryIO
+
+import pytest
+
+from dissect.cstruct import cstruct
+from dissect.cstruct.types import BaseType, MetaType
+
+
+class EtwPointer(BaseType):
+    type: MetaType
+    size: int | None
+
+    @classmethod
+    def _read(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> BaseType:
+        return cls.type._read(stream, context)
+
+    @classmethod
+    def _read_0(cls, stream: BinaryIO, context: dict[str, Any] | None = None) -> list[BaseType]:
+        return cls.type._read_0(stream, context)
+
+    @classmethod
+    def _write(cls, stream: BinaryIO, data: Any) -> int:
+        return cls.type._write(stream, data)
+
+    @classmethod
+    def as_32bit(cls) -> None:
+        cls.type = cls.cs.uint32
+        cls.size = 4
+
+    @classmethod
+    def as_64bit(cls) -> None:
+        cls.type = cls.cs.uint64
+        cls.size = 8
+
+
+def test_adding_custom_type(cs: cstruct) -> None:
+    cs.add_custom_type("EtwPointer", EtwPointer)
+
+    cs.EtwPointer.as_64bit()
+    assert cs.EtwPointer.type is cs.uint64
+    assert len(cs.EtwPointer) == 8
+    assert cs.EtwPointer(b"\xDE\xAD\xBE\xEF" * 2).dumps() == b"\xDE\xAD\xBE\xEF" * 2
+
+    cs.EtwPointer.as_32bit()
+    assert cs.EtwPointer.type is cs.uint32
+    assert len(cs.EtwPointer) == 4
+    assert cs.EtwPointer(b"\xDE\xAD\xBE\xEF" * 2).dumps() == b"\xDE\xAD\xBE\xEF"
+
+
+def test_using_type_in_struct(cs: cstruct) -> None:
+    cs.add_custom_type("EtwPointer", EtwPointer)
+
+    struct_definition = """
+    struct test {
+        EtwPointer data;
+        uint64     data2;
+    };
+    """
+
+    cs.load(struct_definition)
+
+    cs.EtwPointer.as_64bit()
+    assert len(cs.test().data) == 8
+
+    with pytest.raises(EOFError):
+        # Input too small
+        cs.test(b"\xDE\xAD\xBE\xEF" * 3)
+
+    cs.EtwPointer.as_32bit()
+    assert len(cs.test().data) == 4
+
+    assert cs.test(b"\xDE\xAD\xBE\xEF" * 3).data.dumps() == b"\xDE\xAD\xBE\xEF"