Merge pull request #1252 from serpilliere/api_dis_callback

Updt dis_block_callback; apply_splitting

5 files changed, 81 insertions, 90 deletions

diff --git a/example/disasm/callback.py b/example/disasm/callback.py
index 7219462f..1498b11e 100644
--- a/example/disasm/callback.py
+++ b/example/disasm/callback.py
@@ -4,7 +4,7 @@ from miasm.analysis.machine import Machine
 from miasm.core.asmblock import AsmConstraint
 
 
-def cb_x86_callpop(cur_bloc, loc_db, *args, **kwargs):
+def cb_x86_callpop(mdis, cur_bloc, offset_to_dis):
     """
     1000: call 1005
     1005: pop
@@ -28,7 +28,7 @@ def cb_x86_callpop(cur_bloc, loc_db, *args, **kwargs):
         return
 
     loc_key = dst.loc_key
-    offset = loc_db.get_location_offset(loc_key)
+    offset = mdis.loc_db.get_location_offset(loc_key)
     ## The destination must be the next instruction
     if offset != last_instr.offset + last_instr.l:
         return
diff --git a/miasm/arch/arm/disasm.py b/miasm/arch/arm/disasm.py
index 4c92bf6a..e639f327 100644
--- a/miasm/arch/arm/disasm.py
+++ b/miasm/arch/arm/disasm.py
@@ -4,7 +4,7 @@ from miasm.core.asmblock import AsmConstraint, disasmEngine
 from miasm.arch.arm.arch import mn_arm, mn_armt
 
 
-def cb_arm_fix_call(mn, cur_bloc, loc_db, offsets_to_dis, *args, **kwargs):
+def cb_arm_fix_call(mdis, cur_block, offsets_to_dis):
     """
     for arm:
     MOV        LR, PC
@@ -12,22 +12,22 @@ def cb_arm_fix_call(mn, cur_bloc, loc_db, offsets_to_dis, *args, **kwargs):
     * is a subcall *
 
     """
-    if len(cur_bloc.lines) < 2:
+    if len(cur_block.lines) < 2:
         return
-    l1 = cur_bloc.lines[-1]
-    l2 = cur_bloc.lines[-2]
+    l1 = cur_block.lines[-1]
+    l2 = cur_block.lines[-2]
     if l1.name != "LDR":
         return
     if l2.name != "MOV":
         return
 
-    values = viewvalues(mn.pc)
+    values = viewvalues(mdis.mn.pc)
     if not l1.args[0] in values:
         return
     if not l2.args[1] in values:
         return
-    loc_key_cst = loc_db.get_or_create_offset_location(l1.offset + 4)
-    cur_bloc.add_cst(loc_key_cst, AsmConstraint.c_next)
+    loc_key_cst = mdis.loc_db.get_or_create_offset_location(l1.offset + 4)
+    cur_block.add_cst(loc_key_cst, AsmConstraint.c_next)
     offsets_to_dis.add(l1.offset + 4)
 
 cb_arm_funcs = [cb_arm_fix_call]
diff --git a/miasm/arch/x86/disasm.py b/miasm/arch/x86/disasm.py
index 01147970..49b7158a 100644
--- a/miasm/arch/x86/disasm.py
+++ b/miasm/arch/x86/disasm.py
@@ -5,9 +5,9 @@ from miasm.arch.x86.arch import mn_x86
 cb_x86_funcs = []
 
 
-def cb_x86_disasm(*args, **kwargs):
+def cb_x86_disasm(mdis, cur_block, offset_to_dis):
     for func in cb_x86_funcs:
-        func(*args, **kwargs)
+        func(mdis, cur_block, offset_to_dis)
 
 
 class dis_x86(disasmEngine):
diff --git a/miasm/core/asmblock.py b/miasm/core/asmblock.py
index 5ac1312f..f412de86 100644
--- a/miasm/core/asmblock.py
+++ b/miasm/core/asmblock.py
@@ -679,73 +679,8 @@ class AsmCFG(DiGraph):
             log_asmblock.info("size: %d max: %d", block.size, block.max_size)
 
     def apply_splitting(self, loc_db, dis_block_callback=None, **kwargs):
-        """Consider @self' bto destinations and split block in @self if one of
-        these destinations jumps in the middle of this block.
-        In order to work, they must be only one block in @self per loc_key in
-        @loc_db (which is true if @self come from the same disasmEngine).
-
-        @loc_db: LocationDB instance associated with @self'loc_keys
-        @dis_block_callback: (optional) if set, this callback will be called on
-        new block destinations
-        @kwargs: (optional) named arguments to pass to dis_block_callback
-        """
-        # Get all possible destinations not yet resolved, with a resolved
-        # offset
-        block_dst = []
-        for loc_key in self.pendings:
-            offset = loc_db.get_location_offset(loc_key)
-            if offset is not None:
-                block_dst.append(offset)
-
-        todo = set(self.blocks)
-        rebuild_needed = False
-
-        while todo:
-            # Find a block with a destination inside another one
-            cur_block = todo.pop()
-            range_start, range_stop = cur_block.get_range()
-
-            for off in block_dst:
-                if not (off > range_start and off < range_stop):
-                    continue
-
-                # `cur_block` must be split at offset `off`from miasm.core.locationdb import LocationDB
-
-                new_b = cur_block.split(loc_db, off)
-                log_asmblock.debug("Split block %x", off)
-                if new_b is None:
-                    log_asmblock.error("Cannot split %x!!", off)
-                    continue
-
-                # Remove pending from cur_block
-                # Links from new_b will be generated in rebuild_edges
-                for dst in new_b.bto:
-                    if dst.loc_key not in self.pendings:
-                        continue
-                    self.pendings[dst.loc_key] = set(pending for pending in self.pendings[dst.loc_key]
-                                                     if pending.waiter != cur_block)
-
-                # The new block destinations may need to be disassembled
-                if dis_block_callback:
-                    offsets_to_dis = set(
-                        self.loc_db.get_location_offset(constraint.loc_key)
-                        for constraint in new_b.bto
-                    )
-                    dis_block_callback(cur_bloc=new_b,
-                                       offsets_to_dis=offsets_to_dis,
-                                       loc_db=loc_db, **kwargs)
-
-                # Update structure
-                rebuild_needed = True
-                self.add_block(new_b)
-
-                # The new block must be considered
-                todo.add(new_b)
-                range_start, range_stop = cur_block.get_range()
-
-        # Rebuild edges to match new blocks'bto
-        if rebuild_needed:
-            self.rebuild_edges()
+        warnings.warn('DEPRECATION WARNING: apply_splitting is member of disasm_engine')
+        raise RuntimeError("Moved api")
 
     def __str__(self):
         out = []
@@ -1261,8 +1196,7 @@ class disasmEngine(object):
      - lines_wd: maximum block's size (in number of instruction)
      - blocs_wd: maximum number of distinct disassembled block
 
-    + callback(arch, attrib, pool_bin, cur_bloc, offsets_to_dis,
-               loc_db)
+    + callback(mdis, cur_block, offsets_to_dis)
      - dis_block_callback: callback after each new disassembled block
     """
 
@@ -1426,10 +1360,7 @@ class disasmEngine(object):
         cur_block.fix_constraints()
 
         if self.dis_block_callback is not None:
-            self.dis_block_callback(mn=self.arch, attrib=self.attrib,
-                                    pool_bin=self.bin_stream, cur_bloc=cur_block,
-                                    offsets_to_dis=offsets_to_dis,
-                                    loc_db=self.loc_db)
+            self.dis_block_callback(self, cur_block, offsets_to_dis)
         return cur_block, offsets_to_dis
 
     def dis_block(self, offset):
@@ -1470,12 +1401,72 @@ class disasmEngine(object):
             todo += nexts
             blocks.add_block(cur_block)
 
-        blocks.apply_splitting(self.loc_db,
-                               dis_block_callback=self.dis_block_callback,
-                               mn=self.arch, attrib=self.attrib,
-                               pool_bin=self.bin_stream)
+        self.apply_splitting(blocks)
         return blocks
 
+    def apply_splitting(self, blocks):
+        """Consider @blocks' bto destinations and split block in @blocks if one
+        of these destinations jumps in the middle of this block.  In order to
+        work, they must be only one block in @self per loc_key in
+
+        @blocks: Asmcfg
+        """
+        # Get all possible destinations not yet resolved, with a resolved
+        # offset
+        block_dst = []
+        for loc_key in blocks.pendings:
+            offset = self.loc_db.get_location_offset(loc_key)
+            if offset is not None:
+                block_dst.append(offset)
+
+        todo = set(blocks.blocks)
+        rebuild_needed = False
+
+        while todo:
+            # Find a block with a destination inside another one
+            cur_block = todo.pop()
+            range_start, range_stop = cur_block.get_range()
+
+            for off in block_dst:
+                if not (off > range_start and off < range_stop):
+                    continue
+
+                # `cur_block` must be split at offset `off`from miasm.core.locationdb import LocationDB
+
+                new_b = cur_block.split(self.loc_db, off)
+                log_asmblock.debug("Split block %x", off)
+                if new_b is None:
+                    log_asmblock.error("Cannot split %x!!", off)
+                    continue
+
+                # Remove pending from cur_block
+                # Links from new_b will be generated in rebuild_edges
+                for dst in new_b.bto:
+                    if dst.loc_key not in blocks.pendings:
+                        continue
+                    blocks.pendings[dst.loc_key] = set(pending for pending in blocks.pendings[dst.loc_key]
+                                                     if pending.waiter != cur_block)
+
+                # The new block destinations may need to be disassembled
+                if self.dis_block_callback:
+                    offsets_to_dis = set(
+                        self.loc_db.get_location_offset(constraint.loc_key)
+                        for constraint in new_b.bto
+                    )
+                    self.dis_block_callback(self, new_b, offsets_to_dis)
+
+                # Update structure
+                rebuild_needed = True
+                blocks.add_block(new_b)
+
+                # The new block must be considered
+                todo.add(new_b)
+                range_start, range_stop = cur_block.get_range()
+
+        # Rebuild edges to match new blocks'bto
+        if rebuild_needed:
+            blocks.rebuild_edges()
+
     def dis_instr(self, offset):
         """Disassemble one instruction at offset @offset and return the
         corresponding instruction instance
diff --git a/test/core/asmblock.py b/test/core/asmblock.py
index e5ccf252..47f92569 100644
--- a/test/core/asmblock.py
+++ b/test/core/asmblock.py
@@ -274,7 +274,7 @@ assert asmcfg.successors(tob.loc_key) == [tob.loc_key]
 # Check split_block
 ## Without condition for a split, no change
 asmcfg_bef = asmcfg.copy()
-asmcfg.apply_splitting(mdis.loc_db)
+mdis.apply_splitting(asmcfg)
 assert asmcfg_bef == asmcfg
 open("graph5.dot", "w").write(asmcfg.dot())
 ## Create conditions for a block split
@@ -283,7 +283,7 @@ tob.bto.add(AsmConstraintTo(inside_firstbbl))
 asmcfg.rebuild_edges()
 assert len(asmcfg.pendings) == 1
 assert inside_firstbbl in asmcfg.pendings
-asmcfg.apply_splitting(mdis.loc_db)
+mdis.apply_splitting(asmcfg)
 ## Check result
 assert len(asmcfg) == 6
 assert len(asmcfg.pendings) == 0