Compatibility with cstruct v4 (#56)

Requires a minimal version of dissect.cstruct version 4.2

dissect/cobaltstrike/beacon.py

@@ -486,8 +486,8 @@ def parse_transform_binary(program: bytes, build: str = "metadata") -> List[Tupl
         value = u32be(d)
         if len(d) != 4 or value == 0:
             break
-        name = TransformStep.reverse.get(value, "")
-        step = getattr(TransformStep, name, None)
+        step = TransformStep(value)
+        name = step.name
         if step is None:
             raise IndexError("Unknown transform step for value: {}".format(value))
         elif step == TransformStep.BUILD:

dissect/cobaltstrike/c2.py

@@ -174,7 +174,7 @@ def enable_reprlib_c2():
 def c2packet_to_record(c2packet: C2Packet) -> Record:
     """Convert `c2packet` to a flow.record."""
     fields = [("bytes", "raw_http")]
-    kv = dict(c2packet._values)
+    kv = c2packet.__dict__
     for field in c2packet._type.fields:
         ftype = str(field.type)
         if ftype.startswith("char"):

dissect/cobaltstrike/c_c2.py

@@ -192,68 +192,6 @@ def typedef_for_enum(enum_class: IntEnum, int_type: str = "uint32") -> str:
 c2struct.load(typedef_for_enum(BeaconCallback))
 c2struct.load(C2_DEF)
 
-
-# Some wrapper classes for some dissect.cstruct structs, mainly so we can use `isinstance()`
-class BeaconMetadata(cstruct.Instance):
-    magic: int
-    size: int
-    aes_rand: bytes
-    ansi_cp: int
-    oem_cp: int
-    bid: int
-    pid: int
-    port: int
-    flag: int
-    ver_major: int
-    ver_minor: int
-    ver_build: int
-    ptr_x64: int
-    ptr_gmh: int
-    ptr_gpa: int
-    ip: int
-    info: bytes
-
-    def __init__(self, *args, **kwargs):
-        instance = c2struct.BeaconMetadata(*args, **kwargs)
-        super().__init__(instance._type, instance._values, instance._sizes)
-
-    def __eq__(self, other):
-        return self._values == other._values
-
-    def __hash__(self):
-        return hash(tuple(self._values.items()))
-
-
-class CallbackPacket(cstruct.Instance):
-    counter: int
-    size: int
-    callback: BeaconCallback
-    data: bytes
-
-    def __init__(self, *args, **kwargs):
-        instance = c2struct.CallbackPacket(*args, **kwargs)
-        super().__init__(instance._type, instance._values, instance._sizes)
-
-    def __eq__(self, other):
-        return self._values == other._values
-
-    def __hash__(self):
-        return hash(tuple(self._values.items()))
-
-
-class TaskPacket(cstruct.Instance):
-    epoch: int
-    total_size: int
-    command: BeaconCommand
-    size: int
-    data: bytes
-
-    def __init__(self, *args, **kwargs):
-        instance = c2struct.TaskPacket(*args, **kwargs)
-        super().__init__(instance._type, instance._values, instance._sizes)
-
-    def __eq__(self, other):
-        return self._values == other._values
-
-    def __hash__(self):
-        return hash(tuple(self._values.items()))
+BeaconMetadata = c2struct.BeaconMetadata
+CallbackPacket = c2struct.CallbackPacket
+TaskPacket = c2struct.TaskPacket

dissect/cobaltstrike/client.py

@@ -7,6 +7,7 @@
    A default client will perform check-ins and only log the tasks it receives unless implemented otherwise.
 
 """
+
 # Python imports
 import argparse
 import datetime
@@ -267,20 +268,20 @@ def run(
 
         ver_major, ver_minor, ver_build = random_windows_ver()
 
-        self.metadata = BeaconMetadata(
-            magic=0xBEEF,
-            ansi_cp=ansi_cp,
-            oem_cp=oem_cp,
-            bid=self.beacon_id,
-            pid=self.pid,
-            flag=flag,
-            aes_rand=self.aes_rand,
-            ip=internal_ip_int,
-            ver_major=ver_major,
-            ver_minor=ver_minor,
-            ver_build=ver_build,
-            info=info.encode(),
-        )
+        self.metadata = BeaconMetadata()
+        self.metadata.magic = 0xBEEF
+        self.metadata.ansi_cp = ansi_cp
+        self.metadata.oem_cp = oem_cp
+        self.metadata.bid = self.beacon_id
+        self.metadata.pid = self.pid
+        self.metadata.flag = flag
+        self.metadata.aes_rand = self.aes_rand
+        self.metadata.ip = internal_ip_int
+        self.metadata.ver_major = ver_major
+        self.metadata.ver_minor = ver_minor
+        self.metadata.ver_build = ver_build
+        self.metadata.info = info.encode()
+
         self.c2http = C2Http(bconfig, aes_key=self.aes_key, hmac_key=self.hmac_key)
 
         self.domain = domain or random.choice(self.bconfig.domains)

dissect/cobaltstrike/utils.py

@@ -242,7 +242,7 @@ def enable_reprlib_cstruct():
     from dissect.cstruct.types.instance import Instance
 
     def reprlib_repr(self) -> str:
-        values = ", ".join(f"{k}={hex(v) if isinstance(v, int) else reprlib.repr(v)}" for k, v in self._values.items())
+        values = ", ".join(f"{k}={hex(v) if isinstance(v, int) else reprlib.repr(v)}" for k, v in self.__dict__.items())
         return f"<{self._type.name} {values}>"
 
     Instance.__repr__ = reprlib_repr

setup.cfg

@@ -40,7 +40,7 @@ namespace_packages = dissect
 platforms = any
 include_package_data = true
 install_requires =
-    dissect.cstruct >= 2.0, < 4.0
+    dissect.cstruct >= 4.2
     lark
 python_requires = >=3.9
 setup_requires =

tests/test_beacon.py

@@ -212,7 +212,7 @@ def test_beacon_settings_readonly(beacon_x64_file):
             "beacon_x86_file",
             ["-t", "raw"],
             0,
-            b"<Setting index=<BeaconSetting.SETTING_WATERMARK: 37>, type=<SettingsType.TYPE_INT: 2>",
+            b"<Setting index=<BeaconSetting.SETTING_WATERMARK: 37> type=<SettingsType.TYPE_INT: 2>",
             None,
             id="beacon_x86_file-stdin-raw",
         ),

tests/test_c2.py

@@ -374,7 +374,7 @@ def test_beaconmetadata_set_membership():
     m2 = BeaconMetadata(magic=0xBEEF, pid=0x12345, info=b"testing")
     m3 = BeaconMetadata(magic=0xBEEF, pid=0x12345, info=b"12345")
     metadata_set.add(m1)
-    assert list(m1._values.items()) == list(m2._values.items())
+    assert list(m1.__dict__.items()) == list(m2.__dict__.items())
     assert hash(m1) == hash(m2)
     assert m1 == m2
     assert m2 in metadata_set