From 56fdd6bb9713a56e03de8871bffeb9e0323e2745 Mon Sep 17 00:00:00 2001
From: Yun Zheng Hu <hu@fox-it.com>
Date: Mon, 14 Oct 2024 14:28:27 +0200
Subject: [PATCH] Output bof_reuse_memory and bof_allocator in c2profile (#67)

Also the correct allocator name is displayed when pretty printing SETTING_BOF_ALLOCATOR.
---
 dissect/cobaltstrike/beacon.py    |  8 ++++++
 dissect/cobaltstrike/c2profile.py |  4 +++
 tests/test_c2profile.py           | 44 +++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/dissect/cobaltstrike/beacon.py b/dissect/cobaltstrike/beacon.py
index 5618c1b..0ac690c 100644
--- a/dissect/cobaltstrike/beacon.py
+++ b/dissect/cobaltstrike/beacon.py
@@ -247,6 +247,12 @@
     CreateRemoteThread_ = 7,
     NtQueueApcThread_s = 8
 };
+
+enum BofAllocator: uint16 {
+    VirtualAlloc = 0,
+    MapViewOfFile = 1,
+    HeapAlloc = 2,
+};
 """
 
 cs_struct = cstruct.cstruct(endian=">")
@@ -262,6 +268,7 @@
 ProxyServer = cs_struct.ProxyServer
 InjectAllocator = cs_struct.InjectAllocator
 InjectExecutor = cs_struct.InjectExecutor
+BofAllocator = cs_struct.BofAllocator
 
 DEFAULT_XOR_KEYS: List[bytes] = [b"\x69", b"\x2e", b"\x00"]
 """ Default XOR keys used by Cobalt Strike for obfuscating Beacon config bytes """
@@ -635,6 +642,7 @@ def null_terminated_str(data: bytes) -> str:
     BeaconSetting.SETTING_DNS_IDLE: lambda x: str(ipaddress.IPv4Address(x)),
     BeaconSetting.SETTING_WATERMARKHASH: lambda x: null_terminated_bytes(x) if isinstance(x, bytes) else x,
     BeaconSetting.SETTING_MASKED_WATERMARK: lambda x: x.hex(),
+    BeaconSetting.SETTING_BOF_ALLOCATOR: lambda x: BofAllocator(x).name,
     # BeaconSetting.SETTING_PROTOCOL: lambda x: BeaconProtocol(x).name,
     # BeaconSetting.SETTING_CRYPTO_SCHEME: lambda x: CryptoScheme(x).name,
     # BeaconSetting.SETTING_PROXY_BEHAVIOR: lambda x: ProxyServer(x).name,
diff --git a/dissect/cobaltstrike/c2profile.py b/dissect/cobaltstrike/c2profile.py
index d820446..d1b7d77 100644
--- a/dissect/cobaltstrike/c2profile.py
+++ b/dissect/cobaltstrike/c2profile.py
@@ -618,6 +618,10 @@ def from_beacon_config(cls, config: BeaconConfig) -> "C2Profile":
                 dns_beacon.set_option("dns_sleep", value)
             elif setting == BeaconSetting.SETTING_MAXDNS:
                 dns_beacon.set_option("maxdns", value)
+            elif setting == BeaconSetting.SETTING_PROCINJ_BOF_REUSE_MEM and value:
+                proc_inj.set_option("bof_reuse_memory", "true")
+            elif setting == BeaconSetting.SETTING_BOF_ALLOCATOR:
+                proc_inj.set_option("bof_allocator", value)
 
         if c2_recover:
             http_get.set_non_empty_config_block("server", HttpOptionsBlock(output=DataTransformBlock(steps=c2_recover)))
diff --git a/tests/test_c2profile.py b/tests/test_c2profile.py
index f3429c0..29e4615 100644
--- a/tests/test_c2profile.py
+++ b/tests/test_c2profile.py
@@ -1,3 +1,5 @@
+import pytest
+
 from dissect.cobaltstrike import beacon, c2profile
 
 C2PROFILE_SOURCE = """
@@ -294,3 +296,45 @@ def test_c2profile_v48():
     """
     profile = c2profile.C2Profile.from_text(c2profile_fragment)
     assert profile.properties["stage.syscall_method"] == ["direct"]
+
+
+@pytest.mark.parametrize(
+    ("allocator_enum", "bof_allocator"),
+    [
+        (beacon.BofAllocator.HeapAlloc, "HeapAlloc"),
+        (beacon.BofAllocator.MapViewOfFile, "MapViewOfFile"),
+        (beacon.BofAllocator.HeapAlloc, "HeapAlloc"),
+    ],
+)
+def test_c2profile_bof_allocator(allocator_enum, bof_allocator):
+    data = beacon.Setting(
+        index=beacon.BeaconSetting.SETTING_BOF_ALLOCATOR,
+        type=beacon.SettingsType.TYPE_SHORT,
+        length=0x2,
+        value=allocator_enum.dumps(),
+    ).dumps()
+    bconfig = beacon.BeaconConfig(data)
+    profile = c2profile.C2Profile.from_beacon_config(bconfig)
+    assert profile.properties["process-inject.bof_allocator"] == [bof_allocator]
+
+
+@pytest.mark.parametrize(
+    ("bof_reuse_memory",),
+    [
+        (True,),
+        (False,),
+    ],
+)
+def test_c2profile_bof_reuse_memory(bof_reuse_memory):
+    data = beacon.Setting(
+        index=beacon.BeaconSetting.SETTING_PROCINJ_BOF_REUSE_MEM,
+        type=beacon.SettingsType.TYPE_SHORT,
+        length=0x2,
+        value=beacon.cs_struct.uint16(bof_reuse_memory).dumps(),
+    ).dumps()
+    bconfig = beacon.BeaconConfig(data)
+    profile = c2profile.C2Profile.from_beacon_config(bconfig)
+    if bof_reuse_memory:
+        assert profile.properties["process-inject.bof_reuse_memory"] == ["true"]
+    else:
+        assert "project.inject.bof_reuse_memory" not in profile.properties