Example/IDA: add a script to highlight depnodes through IDA

1 files changed, 201 insertions, 0 deletions

diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
new file mode 100644
index 00000000..65a3c88f
--- /dev/null
+++ b/example/ida/depgraph.py
@@ -0,0 +1,201 @@
+import sys
+
+from idaapi import GraphViewer
+from miasm2.core.bin_stream_ida import bin_stream_ida
+from miasm2.core.asmbloc import *
+from miasm2.expression.simplifications import expr_simp
+from miasm2.analysis.machine import Machine
+from miasm2.analysis.depgraph import DependencyGraph, DependencyGraph_NoMemory
+
+from utils import guess_machine
+
+
+class depGraphSettingsForm(Form):
+
+    def __init__(self, ira):
+
+        self.ira = ira
+        self.address = ScreenEA()
+        cur_bloc = list(ira.getby_offset(self.address))[0]
+        for line_nb, l in enumerate(cur_bloc.lines):
+            if l.offset == self.address:
+                break
+        cur_label = str(cur_bloc.label)
+        labels = sorted(map(str, ira.blocs.keys()))
+        regs = sorted(ir_arch.arch.regs.all_regs_ids_byname.keys())
+        reg_default = regs[0]
+        for i in xrange(10):
+            opnd = GetOpnd(self.address, i).upper()
+            if opnd in regs:
+                reg_default = opnd
+                break
+
+        Form.__init__(self,
+r"""BUTTON YES* Launch
+BUTTON CANCEL NONE
+Dependency Graph Settings
+
+Track the element:
+<Before the line:{rBeforeLine}>
+<After the line:{rAfterLine}>
+<At the end of the basic block:{rEndBlock}>{cMode}>
+
+<Target basic block:{cbBBL}>
+<Register to track:{cbReg}>
+<Line number:{iLineNb}>
+
+Method to use:
+<Best effort:{rBest}>
+<No memory (sound & complete):{rNoMem}>{cMethod}>
+
+<Highlight color:{cColor}>
+""", {
+            'cbReg': Form.DropdownListControl(
+                    items=regs,
+                    readonly=False,
+                    selval=reg_default),
+            'cMode': Form.RadGroupControl(("rBeforeLine", "rAfterLine",
+                                           "rEndBlock")),
+            'cMethod': Form.RadGroupControl(("rBest", "rNoMem")),
+            'iLineNb': Form.NumericInput(tp=Form.FT_RAWHEX,
+                                         value=line_nb),
+            'cbBBL': Form.DropdownListControl(
+                    items=labels,
+                    readonly=False,
+                    selval=cur_label),
+            'cColor': Form.ColorInput(value=0xc0c020),
+        })
+
+        self.Compile()
+
+    @property
+    def label(self):
+        value = self.cbBBL.value
+        for real_label in self.ira.blocs:
+            if str(real_label) == value:
+                return real_label
+        raise ValueError("Bad label")
+
+    @property
+    def line_nb(self):
+        value = self.iLineNb.value
+        mode = self.cMode.value
+        if mode == 0:
+            return value
+        elif mode == 1:
+            return value + 1
+        else:
+            return len(self.ira.blocs[self.label].irs)
+
+    @property
+    def elements(self):
+        value = self.cbReg.value
+        return set([ir_arch.arch.regs.all_regs_ids_byname[value]])
+
+    @property
+    def method(self):
+        value = self.cMethod.value
+        if value == 0:
+            return DependencyGraph
+        elif value == 1:
+            return DependencyGraph_NoMemory
+        else:
+            raise ValueError("Unknown method")
+
+    @property
+    def color(self):
+        return self.cColor.value
+
+
+# Init
+machine = guess_machine()
+mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira
+
+bs = bin_stream_ida()
+mdis = dis_engine(bs, dont_dis_nulstart_bloc=True)
+ir_arch = ira(mdis.symbol_pool)
+
+# Populate symbols with ida names
+for ad, name in Names():
+    if name is None:
+        continue
+    mdis.symbol_pool.add_label(name, ad)
+
+# Get the current function
+addr = ScreenEA()
+func = idaapi.get_func(addr)
+blocs = mdis.dis_multibloc(func.startEA)
+
+# Generate IR
+for bloc in blocs:
+    ir_arch.add_bloc(bloc)
+
+# Simplify affectations
+for irb in ir_arch.blocs.values():
+    for irs in irb.irs:
+        for i, e in enumerate(irs):
+            e.dst, e.src = expr_simp(e.dst), expr_simp(e.src)
+
+# Build the IRA Graph
+ir_arch.gen_graph()
+
+# Get settings
+settings = depGraphSettingsForm(ir_arch)
+settings.Execute()
+
+# Get dependency graphs
+dg = (settings.method)(ir_arch)
+graphs = dg.get(settings.label, settings.elements, settings.line_nb,
+                set([ir_arch.symbol_pool.getby_offset(func.startEA)]))
+
+# Display the result
+comments = {}
+sol_nb = 0
+
+def clean_lines():
+    "Remove previous comments"
+    global comments
+    for offset in comments:
+        SetColor(offset, CIC_ITEM, 0xffffff)
+        MakeComm(offset, "")
+    comments = {}
+
+def treat_element():
+    "Display an element"
+    global graphs, comments, sol_nb, settings
+
+    try:
+        graph = graphs.next()
+    except StopIteration:
+        comments = {}
+        print "Done: %d solutions" % (sol_nb)
+        return
+
+    sol_nb += 1
+    print "Get graph number %02d" % sol_nb
+    filename = "/tmp/solution_0x%08x_%02d.dot" % (addr, sol_nb)
+    print "Dump the graph to %s" % filename
+    open(filename, "w").write(graph.graph.dot())
+
+    for node in graph.relevant_nodes:
+        try:
+            offset = ir_arch.blocs[node.label].lines[node.line_nb].offset
+        except IndexError:
+            print "Unable to highlight %s" % node
+            continue
+        comments[offset] = comments.get(offset, []) + [node.element]
+        SetColor(offset, CIC_ITEM, settings.color)
+
+    print "Possible value: %s" % graph.emul().values()[0]
+
+    for offset, elements in comments.iteritems():
+        MakeComm(offset, ", ".join(map(str, elements)))
+
+def next_element():
+    "Display the next element"
+    clean_lines()
+    treat_element()
+
+# Register and launch
+idaapi.add_hotkey("Shift-N", next_element)
+treat_element()