Added api.Hook.get_arguments

Author: ohjeongwook

shellcode_emulator/api.py

@@ -16,8 +16,10 @@
 logger = logging.getLogger(__name__)
 
 class Hook:
-    def __init__(self, emulator):
+    amd64_argument_regs = (UC_X86_REG_RCX, UC_X86_REG_RDX, UC_X86_REG_R8, UC_X86_REG_R9)
+    def __init__(self, emulator, arch):
         self.emulator = emulator
+        self.arch = arch
         self.uc = emulator.uc
         self.trace_target_modules = ['ntdll', 'kernel32', 'kernelbase']
         self.last_code_information = {}
@@ -31,20 +33,34 @@ def return_function(self, uc, return_address, arg_count, return_value):
         uc.reg_write(self.emulator.register.get_by_name("sp"), esp+4*(arg_count+1))        
         uc.reg_write(self.emulator.register.get_by_name("ax"), return_value)
 
-    def callback(self, uc, address, size, user_data):
-        self.emulator.instruction.dump_disassembly(address, size, find_symbol = True)
+    def get_arguments(self, count):
+        arguments = []
+
+        stack_argument_count = count
+        if self.arch == 'AMD64':
+            reg_argument_count = min(count, len(self.amd64_argument_regs))
+            stack_argument_count -= reg_argument_count
+            for i in range(0, reg_argument_count, 1):
+                arguments.append(self.uc.reg_read(self.amd64_argument_regs[i]))
+
+        if stack_argument_count > 0:
+            arguments += self.emulator.memory.get_stack(stack_argument_count)
 
+        return arguments        
+
+    def callback(self, uc, address, size, user_data):
         code = uc.mem_read(address, size)
         try:
             name = self.emulator.debugger.find_symbol(address)
         except:
             name = ''
+            self.emulator.instruction.dump_disassembly(address, size, find_symbol = True)
 
-        print('callback: %s (%x)' % (name, address))
+        sp = self.uc.reg_read(self.emulator.register.get_by_name("sp"))
+        print('%x: %s (%x)' % (sp, name, address))
         if name == 'kernel32!WinExec':
-            rcx = uc.reg_read(UC_X86_REG_RCX)
-            print('rcx: %x' % rcx)
-            command = self.emulator.memory.read_string(rcx)
+            arguments = self.get_arguments(2)
+            command = self.emulator.memory.read_string(arguments[0])
             print('command: %s' % command)
 
         elif name == 'ntdll!LdrLoadDll':

shellcode_emulator/run.py

@@ -53,6 +53,7 @@ def start(self, start, end):
 
 class ShellcodeEmulator:
     def __init__(self, shellcode_filename, shellcode_bytes = '', dump_filename = '', arch = 'AMD64', exhaustive_loop_dump_frequency = 0x10000):
+        self.arch = arch
         self.shellcode_filename = shellcode_filename
         self.shellcode_bytes = shellcode_bytes
         self.exhaustive_loop_dump_frequency = exhaustive_loop_dump_frequency
@@ -102,7 +103,7 @@ def run(self, trace_self_modification = False, print_first_instructions = False)
             self.emulator.add_unicorn_hook(UC_HOOK_CODE, self.instruction_callback, None, self.code_start, self.code_start+1)
 
         self.emulator.memory.hook_unmapped_memory_access()
-        api_hook = shellcode_emulator.api.Hook(self.emulator)
+        api_hook = shellcode_emulator.api.Hook(self.emulator, self.arch)
         api_hook.start()
 
         self.emulator.instruction.set_code_range(self.code_start, self.code_start+self.code_length)