diff options
Diffstat (limited to 'example/ida/graph_ir.py')
| -rw-r--r-- | example/ida/graph_ir.py | 161 |
1 files changed, 91 insertions, 70 deletions
diff --git a/example/ida/graph_ir.py b/example/ida/graph_ir.py index 678ae2ff..8c7e80f5 100644 --- a/example/ida/graph_ir.py +++ b/example/ida/graph_ir.py @@ -1,7 +1,9 @@ import os import tempfile -from idaapi import GraphViewer +import idaapi +import idc +import idautils from miasm2.core.bin_stream_ida import bin_stream_ida from miasm2.core.asmblock import expr_is_label, AsmLabel, is_int @@ -34,7 +36,7 @@ def label_str(self): AsmLabel.__init__ = label_init AsmLabel.__str__ = label_str -def color_irblock(irblock): +def color_irblock(irblock, ir_arch): out = [] lbl = idaapi.COLSTR(str(irblock.label), idaapi.SCOLOR_INSN) out.append(lbl) @@ -52,27 +54,25 @@ def color_irblock(irblock): return "\n".join(out) -class GraphMiasmIR(GraphViewer): +class GraphMiasmIR(idaapi.GraphViewer): def __init__(self, ir_arch, title, result): - GraphViewer.__init__(self, title) - print 'init' + idaapi.GraphViewer.__init__(self, title) self.ir_arch = ir_arch self.result = result self.names = {} def OnRefresh(self): - print 'refresh' self.Clear() addr_id = {} for irblock in self.ir_arch.blocks.values(): - id_irblock = self.AddNode(color_irblock(irblock)) + id_irblock = self.AddNode(color_irblock(irblock, self.ir_arch)) addr_id[irblock] = id_irblock for irblock in self.ir_arch.blocks.values(): if not irblock: continue - all_dst = ir_arch.dst_trackback(irblock) + all_dst = self.ir_arch.dst_trackback(irblock) for dst in all_dst: if not expr_is_label(dst): continue @@ -102,7 +102,7 @@ class GraphMiasmIR(GraphViewer): print "command:", cmd_id def Show(self): - if not GraphViewer.Show(self): + if not idaapi.GraphViewer.Show(self): return False self.cmd_test = self.AddCommand("Test", "F2") if self.cmd_test == 0: @@ -110,64 +110,85 @@ class GraphMiasmIR(GraphViewer): return True -machine = guess_machine() -mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira - -print "Arch", dis_engine - -fname = GetInputFile() -print fname - -bs = bin_stream_ida() -mdis = dis_engine(bs) -ir_arch = ira(mdis.symbol_pool) - -# populate symbols with ida names -for addr, name in Names(): - # print hex(ad), repr(name) - if name is None: - continue - mdis.symbol_pool.add_label(name, addr) - -print "start disasm" -addr = ScreenEA() -print hex(addr) - -blocks = mdis.dis_multibloc(addr) - -print "generating graph" -open('asm_flow.dot', 'w').write(blocks.dot()) - -print "generating IR... %x" % addr - -for block in blocks: - print 'ADD' - print block - ir_arch.add_bloc(block) - - -print "IR ok... %x" % addr - -for irb in ir_arch.blocks.itervalues(): - irs = [] - for assignblk in irb.irs: - new_assignblk = { - expr_simp(dst): expr_simp(src) - for dst, src in assignblk.iteritems() - } - irs.append(AssignBlock(new_assignblk, instr=assignblk.instr)) - ir_arch.blocks[irb.label] = IRBlock(irb.label, irs) - -out = ir_arch.graph.dot() -open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out) - - -# dead_simp(ir_arch) - -g = GraphMiasmIR(ir_arch, "Miasm IR graph", None) - -g.cmd_a = g.AddCommand("cmd a", "x") -g.cmd_b = g.AddCommand("cmd b", "y") - -g.Show() - +def build_graph(verbose=False, simplify=False): + machine = guess_machine() + mn, dis_engine, ira = machine.mn, machine.dis_engine, machine.ira + + if verbose: + print "Arch", dis_engine + + fname = idc.GetInputFile() + if verbose: + print fname + + bs = bin_stream_ida() + mdis = dis_engine(bs) + ir_arch = ira(mdis.symbol_pool) + + # populate symbols with ida names + for addr, name in idautils.Names(): + # print hex(ad), repr(name) + if name is None: + continue + mdis.symbol_pool.add_label(name, addr) + + if verbose: + print "start disasm" + addr = idc.ScreenEA() + if verbose: + print hex(addr) + + blocks = mdis.dis_multibloc(addr) + + if verbose: + print "generating graph" + open('asm_flow.dot', 'w').write(blocks.dot()) + + print "generating IR... %x" % addr + + for block in blocks: + if verbose: + print 'ADD' + print block + ir_arch.add_bloc(block) + + if verbose: + print "IR ok... %x" % addr + + for irb in ir_arch.blocks.itervalues(): + irs = [] + for assignblk in irb.irs: + new_assignblk = { + expr_simp(dst): expr_simp(src) + for dst, src in assignblk.iteritems() + } + irs.append(AssignBlock(new_assignblk, instr=assignblk.instr)) + ir_arch.blocks[irb.label] = IRBlock(irb.label, irs) + + if verbose: + out = ir_arch.graph.dot() + open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out) + title = "Miasm IR graph" + + if simplify: + dead_simp(ir_arch) + + ir_arch.simplify(expr_simp) + modified = True + while modified: + modified = False + modified |= dead_simp(ir_arch) + modified |= ir_arch.remove_empty_assignblks() + modified |= ir_arch.remove_jmp_blocks() + modified |= ir_arch.merge_blocks() + title += " (simplified)" + + g = GraphMiasmIR(ir_arch, title, None) + + g.cmd_a = g.AddCommand("cmd a", "x") + g.cmd_b = g.AddCommand("cmd b", "y") + + g.Show() + +if __name__ == "__main__": + build_graph(verbose=True, simplify=False) |