IDA/Depgraph: clean & fix imports, export main function

1 files changed, 93 insertions, 83 deletions

diff --git a/example/ida/depgraph.py b/example/ida/depgraph.py
index ab033e15..12d01af9 100644
--- a/example/ida/depgraph.py
+++ b/example/ida/depgraph.py
@@ -1,6 +1,11 @@
 import os
 import tempfile
 
+import idautils
+import idc
+import ida_funcs
+import ida_kernwin
+
 from miasm2.core.bin_stream_ida import bin_stream_ida
 from miasm2.core.asmblock import *
 from miasm2.expression import expression as m2_expr
@@ -12,8 +17,7 @@ from miasm2.ir.ir import AssignBlock, IRBlock
 from utils import guess_machine
 
 
-
-class depGraphSettingsForm(Form):
+class depGraphSettingsForm(ida_kernwin.Form):
 
     def __init__(self, ira):
 
@@ -21,7 +25,7 @@ class depGraphSettingsForm(Form):
         self.stk_args = {'ARG%d' % i:i for i in xrange(10)}
         self.stk_unalias_force = False
 
-        self.address = ScreenEA()
+        self.address = idc.ScreenEA()
         cur_block = None
         for block in ira.getby_offset(self.address):
             if block.label.offset is not None:
@@ -36,16 +40,16 @@ class depGraphSettingsForm(Form):
         assert line_nb is not None
         cur_label = str(cur_block.label)
         labels = sorted(map(str, ira.blocks.keys()))
-        regs = sorted(ir_arch.arch.regs.all_regs_ids_byname.keys())
+        regs = sorted(ira.arch.regs.all_regs_ids_byname.keys())
         regs += self.stk_args.keys()
         reg_default = regs[0]
         for i in xrange(10):
-            opnd = GetOpnd(self.address, i).upper()
+            opnd = idc.GetOpnd(self.address, i).upper()
             if opnd in regs:
                 reg_default = opnd
                 break
 
-        Form.__init__(self,
+        ida_kernwin.Form.__init__(self,
 r"""BUTTON YES* Launch
 BUTTON CANCEL NONE
 Dependency Graph Settings
@@ -69,21 +73,22 @@ Method to use:
 
 <Highlight color:{cColor}>
 """, {
-            'cbReg': Form.DropdownListControl(
+            'cbReg': ida_kernwin.Form.DropdownListControl(
                     items=regs,
                     readonly=False,
                     selval=reg_default),
-            'cMode': Form.RadGroupControl(("rBeforeLine", "rAfterLine",
-                                           "rEndBlock")),
-            'cMethod': Form.ChkGroupControl(("rNoMem", "rNoCall", "rImplicit",
-                                             "rUnaliasStack")),
-            'iLineNb': Form.NumericInput(tp=Form.FT_RAWHEX,
-                                         value=line_nb),
-            'cbBBL': Form.DropdownListControl(
+            'cMode': ida_kernwin.Form.RadGroupControl(
+                ("rBeforeLine", "rAfterLine", "rEndBlock")),
+            'cMethod': ida_kernwin.Form.ChkGroupControl(
+                ("rNoMem", "rNoCall", "rImplicit", "rUnaliasStack")),
+            'iLineNb': ida_kernwin.Form.NumericInput(
+                tp=ida_kernwin.Form.FT_RAWHEX,
+                value=line_nb),
+            'cbBBL': ida_kernwin.Form.DropdownListControl(
                     items=labels,
                     readonly=False,
                     selval=cur_label),
-            'cColor': Form.ColorInput(value=0xc0c020),
+            'cColor': ida_kernwin.Form.ColorInput(value=0xc0c020),
         })
 
         self.Compile()
@@ -113,14 +118,14 @@ Method to use:
         if value in self.stk_args:
             line = self.ira.blocks[self.label].irs[self.line_nb].instr
             arg_num = self.stk_args[value]
-            stk_high = m2_expr.ExprInt(GetSpd(line.offset), ir_arch.sp.size)
+            stk_high = m2_expr.ExprInt(idc.GetSpd(line.offset), ir_arch.sp.size)
             stk_off = m2_expr.ExprInt(self.ira.sp.size/8 * arg_num, ir_arch.sp.size)
             element =  m2_expr.ExprMem(mn.regs.regs_init[ir_arch.sp] + stk_high + stk_off, self.ira.sp.size)
             element = expr_simp(element)
             # Force stack unaliasing
             self.stk_unalias_force = True
         elif value:
-            element = ir_arch.arch.regs.all_regs_ids_byname.get(value, None)
+            element = self.ira.arch.regs.all_regs_ids_byname.get(value, None)
 
         else:
             raise ValueError("Unknown element '%s'!" % value)
@@ -142,75 +147,17 @@ Method to use:
     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)
-blocks = mdis.dis_multibloc(func.startEA)
-
-# Generate IR
-for block in blocks:
-    ir_arch.add_bloc(block)
-
-# Get settings
-settings = depGraphSettingsForm(ir_arch)
-settings.Execute()
-
-label, elements, line_nb = settings.label, settings.elements, settings.line_nb
-# Simplify affectations
-for irb in ir_arch.blocks.values():
-    irs = []
-    fix_stack = irb.label.offset is not None and settings.unalias_stack
-    for assignblk in irb.irs:
-        if fix_stack:
-            stk_high = m2_expr.ExprInt(GetSpd(assignblk.instr.offset), ir_arch.sp.size)
-            fix_dct = {ir_arch.sp: mn.regs.regs_init[ir_arch.sp] + stk_high}
-
-        new_assignblk = {}
-        for dst, src in assignblk.iteritems():
-            if fix_stack:
-                src = src.replace_expr(fix_dct)
-                if dst != ir_arch.sp:
-                    dst = dst.replace_expr(fix_dct)
-            dst, src = expr_simp(dst), expr_simp(src)
-            new_assignblk[dst] = src
-        irs.append(AssignBlock(new_assignblk, instr=assignblk.instr))
-    ir_arch.blocks[irb.label] = IRBlock(irb.label, irs)
-
-# Get dependency graphs
-dg = settings.depgraph
-graphs = dg.get(label, elements, 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, "")
+        idc.SetColor(offset, idc.CIC_ITEM, 0xffffff)
+        idc.MakeComm(offset, "")
     comments = {}
 
 def treat_element():
     "Display an element"
-    global graphs, comments, sol_nb, settings
+    global graphs, comments, sol_nb, settings, addr, ir_arch
 
     try:
         graph = graphs.next()
@@ -232,7 +179,7 @@ def treat_element():
             print "Unable to highlight %s" % node
             continue
         comments[offset] = comments.get(offset, []) + [node.element]
-        SetColor(offset, CIC_ITEM, settings.color)
+        idc.SetColor(offset, idc.CIC_ITEM, settings.color)
 
     if graph.has_loop:
         print 'Graph has dependency loop: symbolic execution is inexact'
@@ -240,13 +187,76 @@ def treat_element():
         print "Possible value: %s" % graph.emul().values()[0]
 
     for offset, elements in comments.iteritems():
-        MakeComm(offset, ", ".join(map(str, elements)))
+        idc.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()
+
+def launch_depgraph():
+    global graphs, comments, sol_nb, settings, addr, ir_arch
+    # 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 idautils.Names():
+        if name is None:
+            continue
+        mdis.symbol_pool.add_label(name, ad)
+
+    # Get the current function
+    addr = idc.ScreenEA()
+    func = ida_funcs.get_func(addr)
+    blocks = mdis.dis_multibloc(func.startEA)
+
+    # Generate IR
+    for block in blocks:
+        ir_arch.add_bloc(block)
+
+    # Get settings
+    settings = depGraphSettingsForm(ir_arch)
+    settings.Execute()
+
+    label, elements, line_nb = settings.label, settings.elements, settings.line_nb
+    # Simplify affectations
+    for irb in ir_arch.blocks.values():
+        irs = []
+        fix_stack = irb.label.offset is not None and settings.unalias_stack
+        for assignblk in irb.irs:
+            if fix_stack:
+                stk_high = m2_expr.ExprInt(idc.GetSpd(assignblk.instr.offset), ir_arch.sp.size)
+                fix_dct = {ir_arch.sp: mn.regs.regs_init[ir_arch.sp] + stk_high}
+
+            new_assignblk = {}
+            for dst, src in assignblk.iteritems():
+                if fix_stack:
+                    src = src.replace_expr(fix_dct)
+                    if dst != ir_arch.sp:
+                        dst = dst.replace_expr(fix_dct)
+                dst, src = expr_simp(dst), expr_simp(src)
+                new_assignblk[dst] = src
+            irs.append(AssignBlock(new_assignblk, instr=assignblk.instr))
+        ir_arch.blocks[irb.label] = IRBlock(irb.label, irs)
+
+    # Get dependency graphs
+    dg = settings.depgraph
+    graphs = dg.get(label, elements, line_nb,
+                    set([ir_arch.symbol_pool.getby_offset(func.startEA)]))
+
+    # Display the result
+    comments = {}
+    sol_nb = 0
+
+    # Register and launch
+    ida_kernwin.add_hotkey("Shift-N", next_element)
+    treat_element()
+
+if __name__ == "__main__":
+    launch_depgraph()