about summary refs log tree commit diff stats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--example/disasm/callback.py67
-rw-r--r--miasm2/arch/aarch64/disasm.py4
-rw-r--r--miasm2/arch/arm/disasm.py15
-rw-r--r--miasm2/arch/x86/disasm.py36
-rw-r--r--miasm2/core/asmbloc.py11
-rw-r--r--test/test_all.py1
6 files changed, 87 insertions, 47 deletions
diff --git a/example/disasm/callback.py b/example/disasm/callback.py
new file mode 100644
index 00000000..6c77023e
--- /dev/null
+++ b/example/disasm/callback.py
@@ -0,0 +1,67 @@
+from miasm2.core.bin_stream import bin_stream_str
+from miasm2.core.asmbloc import asm_label, asm_constraint, expr_is_label
+from miasm2.arch.x86.disasm import dis_x86_32, cb_x86_funcs
+
+
+def cb_x86_callpop(cur_bloc, symbol_pool, *args, **kwargs):
+    """
+    1000: call 1005
+    1005: pop
+
+    Will give:
+
+    1000: push 1005
+    1005: pop
+
+    """
+    # Pattern matching
+    if len(cur_bloc.lines) < 1:
+        return
+    ## We want to match a CALL, always the last line of a basic block
+    last_instr = cur_bloc.lines[-1]
+    if last_instr.name != 'CALL':
+        return
+    ## The destination must be a label
+    dst = last_instr.args[0]
+    if not expr_is_label(dst):
+        return
+    ## The destination must be the next instruction
+    if dst.name.offset != last_instr.offset + last_instr.l:
+        return
+
+    # Update instruction instance
+    last_instr.name = 'PUSH'
+
+    # Update next blocks to process in the disassembly engine
+    cur_bloc.bto.clear()
+    cur_bloc.add_cst(dst.name.offset, asm_constraint.c_next, symbol_pool)
+
+
+# Prepare a tiny shellcode
+shellcode = ''.join(["\xe8\x00\x00\x00\x00", # CALL $
+                     "X",                    # POP EAX
+                     "\xc3",                 # RET
+                     ])
+bin_stream = bin_stream_str(shellcode)
+mdis = dis_x86_32(bin_stream)
+
+print "Without callback:\n"
+blocks = mdis.dis_multibloc(0)
+print "\n".join(str(block) for block in blocks)
+
+# Enable callback
+cb_x86_funcs.append(cb_x86_callpop)
+## Other method:
+## mdis.dis_bloc_callback = cb_x86_callpop
+
+# Clean disassembly cache
+mdis.job_done.clear()
+
+print "=" * 40
+print "With callback:\n"
+blocks_after = mdis.dis_multibloc(0)
+print "\n".join(str(block) for block in blocks_after)
+
+# Ensure the callback has been called
+assert blocks[0].lines[0].name == "CALL"
+assert blocks_after[0].lines[0].name == "PUSH"
diff --git a/miasm2/arch/aarch64/disasm.py b/miasm2/arch/aarch64/disasm.py
index 1fc19d07..d31ce3fd 100644
--- a/miasm2/arch/aarch64/disasm.py
+++ b/miasm2/arch/aarch64/disasm.py
@@ -4,9 +4,9 @@ from miasm2.arch.aarch64.arch import mn_aarch64
 cb_aarch64_funcs = []
 
 
-def cb_aarch64_disasm(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool):
+def cb_aarch64_disasm(*args, **kwargs):
     for func in cb_aarch64_funcs:
-        func(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool)
+        func(*args, **kwargs)
 
 
 class dis_aarch64b(disasmEngine):
diff --git a/miasm2/arch/arm/disasm.py b/miasm2/arch/arm/disasm.py
index 03a17562..6209be5e 100644
--- a/miasm2/arch/arm/disasm.py
+++ b/miasm2/arch/arm/disasm.py
@@ -2,8 +2,7 @@ from miasm2.core.asmbloc import asm_constraint, disasmEngine
 from miasm2.arch.arm.arch import mn_arm, mn_armt
 
 
-def cb_arm_fix_call(
-    mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool):
+def cb_arm_fix_call(mn, cur_bloc, symbol_pool, offsets_to_dis, *args, **kwargs):
     """
     for arm:
     MOV        LR, PC
@@ -19,11 +18,11 @@ def cb_arm_fix_call(
         return
     if l2.name != "MOV":
         return
-    # print cur_bloc
-    # print l1
-    if not l1.args[0] in mn.pc.values():
+
+    values = mn.pc.values()
+    if not l1.args[0] in values:
         return
-    if not l2.args[1] in mn.pc.values():
+    if not l2.args[1] in values:
         return
     cur_bloc.add_cst(l1.offset + 4, asm_constraint.c_next, symbol_pool)
     offsets_to_dis.add(l1.offset + 4)
@@ -31,9 +30,9 @@ def cb_arm_fix_call(
 cb_arm_funcs = [cb_arm_fix_call]
 
 
-def cb_arm_disasm(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool):
+def cb_arm_disasm(*args, **kwargs):
     for func in cb_arm_funcs:
-        func(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool)
+        func(*args, **kwargs)
 
 
 class dis_armb(disasmEngine):
diff --git a/miasm2/arch/x86/disasm.py b/miasm2/arch/x86/disasm.py
index 8d06fb8d..0ff55097 100644
--- a/miasm2/arch/x86/disasm.py
+++ b/miasm2/arch/x86/disasm.py
@@ -1,41 +1,13 @@
-from miasm2.core.asmbloc import asm_constraint, asm_label, disasmEngine
-from miasm2.expression.expression import ExprId
+from miasm2.core.asmbloc import disasmEngine
 from miasm2.arch.x86.arch import mn_x86
 
 
-def cb_x86_callpop(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool):
-    """
-    1000: call 1005
-    1005: pop
+cb_x86_funcs = []
 
-    Will give:
 
-    1000: push 1005
-    1005: pop
-
-    """
-
-    if len(cur_bloc.lines) < 1:
-        return
-    l = cur_bloc.lines[-1]
-    if l.name != 'CALL':
-        return
-    dst = l.args[0]
-    if not (isinstance(dst, ExprId) and isinstance(dst.name, asm_label)):
-        return
-    if dst.name.offset != l.offset + l.l:
-        return
-    l.name = 'PUSH'
-    cur_bloc.bto = set()
-    cur_bloc.add_cst(dst.name.offset, asm_constraint.c_next, symbol_pool)
-
-
-cb_x86_funcs = [cb_x86_callpop]
-
-
-def cb_x86_disasm(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool):
+def cb_x86_disasm(*args, **kwargs):
     for func in cb_x86_funcs:
-        func(mn, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool)
+        func(*args, **kwargs)
 
 
 class dis_x86(disasmEngine):
diff --git a/miasm2/core/asmbloc.py b/miasm2/core/asmbloc.py
index 9f67747c..71e577cf 100644
--- a/miasm2/core/asmbloc.py
+++ b/miasm2/core/asmbloc.py
@@ -440,8 +440,9 @@ def dis_bloc(mnemo, pool_bin, cur_bloc, offset, job_done, symbol_pool,
         offsets_to_dis.add(offset)
 
     if dis_bloc_callback is not None:
-        dis_bloc_callback(
-            mnemo, attrib, pool_bin, cur_bloc, offsets_to_dis, symbol_pool)
+        dis_bloc_callback(mn=mnemo, attrib=attrib, pool_bin=pool_bin,
+                          cur_bloc=cur_bloc, offsets_to_dis=offsets_to_dis,
+                          symbol_pool=symbol_pool)
     # print 'dst', [hex(x) for x in offsets_to_dis]
     return offsets_to_dis
 
@@ -482,9 +483,9 @@ def split_bloc(mnemo, attrib, pool_bin, blocs,
                 offsets_to_dis = set(
                     [x.label.offset for x in new_b.bto
                      if isinstance(x.label, asm_label)])
-                dis_bloc_callback(
-                    mnemo, attrib, pool_bin, new_b, offsets_to_dis,
-                    symbol_pool)
+                dis_bloc_callback(mn=mnemo, attrib=attrib, pool_bin=pool_bin,
+                                  cur_bloc=new_b, offsets_to_dis=offsets_to_dis,
+                                  symbol_pool=symbol_pool)
             blocs.append(new_b)
             a, b = cb.get_range()
 
diff --git a/test/test_all.py b/test/test_all.py
index 8d5d16fe..faae79c5 100644
--- a/test/test_all.py
+++ b/test/test_all.py
@@ -250,6 +250,7 @@ class ExampleDisassembler(Example):
 
 
 for script, prods in [(["single_instr.py"], []),
+                      (["callback.py"], []),
                       (["function.py"], ["graph.txt"]),
                       (["file.py", Example.get_sample("box_upx.exe"),
                         "0x407570"], ["graph.txt"]),