diff options
Diffstat (limited to 'example/ida/depgraph.py')
| -rw-r--r-- | example/ida/depgraph.py | 176 |
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() |