Merge pull request #1447 from cea-sec/helper-depgraph

Depgraph: add `.address_to_location` to help find the line and loc of an address

2 files changed, 20 insertions, 7 deletions

diff --git a/example/symbol_exec/depgraph.py b/example/symbol_exec/depgraph.py
index a7df2e25..61dbd4b4 100644
--- a/example/symbol_exec/depgraph.py
+++ b/example/symbol_exec/depgraph.py
@@ -79,16 +79,11 @@ dg = DependencyGraph(
 
 # Build information
 target_addr = int(args.target_addr, 0)
-current_loc_key = next(iter(ircfg.getby_offset(target_addr)))
-assignblk_index = 0
-current_block = ircfg.get_block(current_loc_key)
-for assignblk_index, assignblk in enumerate(current_block):
-    if assignblk.instr.offset == target_addr:
-        break
+target = dg.address_to_location(target_addr)
 
 # Enumerate solutions
 json_solutions = []
-for sol_nb, sol in enumerate(dg.get(current_block.loc_key, elements, assignblk_index, set())):
+for sol_nb, sol in enumerate(dg.get(target["loc_key"], elements, target["line_nb"], set())):
     fname = "sol_%d.dot" % sol_nb
     with open(fname, "w") as fdesc:
             fdesc.write(sol.graph.dot())
diff --git a/miasm/analysis/depgraph.py b/miasm/analysis/depgraph.py
index 7fadd9bf..436e5354 100644
--- a/miasm/analysis/depgraph.py
+++ b/miasm/analysis/depgraph.py
@@ -639,3 +639,21 @@ class DependencyGraph(object):
         lead = list(depnodes)[0]
         elements = set(depnode.element for depnode in depnodes)
         return self.get(lead.loc_key, elements, lead.line_nb, heads)
+
+    def address_to_location(self, address):
+        """Helper to retrieve the .get() arguments, ie.
+        assembly address -> irblock's location key and line number
+        """
+        current_loc_key = next(iter(self._ircfg.getby_offset(address)))
+        assignblk_index = 0
+        current_block = self._ircfg.get_block(current_loc_key)
+        for assignblk_index, assignblk in enumerate(current_block):
+            if assignblk.instr.offset == address:
+                break
+        else:
+            return None
+
+        return {
+            "loc_key": current_block.loc_key,
+            "line_nb": assignblk_index,
+        }