def build_graph(start_addr, type_graph, simplify=False, dontmodstack=True, loadint=False, verbose=False): machine = guess_machine(addr=start_addr) dis_engine, ira = machine.dis_engine, machine.ira class IRADelModCallStack(ira): def call_effects(self, addr, instr): assignblks, extra = super(IRADelModCallStack, self).call_effects(addr, instr) if not dontmodstack: return assignblks, extra out = [] for assignblk in assignblks: dct = dict(assignblk) dct = { dst: src for (dst, src) in dct.iteritems() if dst != self.sp } out.append(AssignBlock(dct, assignblk.instr)) return out, extra if verbose: print "Arch", dis_engine fname = idc.GetInputFile() if verbose: print fname bs = bin_stream_ida() mdis = dis_engine(bs) ir_arch = IRADelModCallStack(mdis.loc_db) # populate symbols with ida names for addr, name in idautils.Names(): if name is None: continue if (mdis.loc_db.get_offset_location(addr) or mdis.loc_db.get_name_location(name)): # Symbol alias continue mdis.loc_db.add_location(name, addr) if verbose: print "start disasm" if verbose: print hex(start_addr) asmcfg = mdis.dis_multiblock(start_addr) entry_points = set([mdis.loc_db.get_offset_location(start_addr)]) if verbose: print "generating graph" open('asm_flow.dot', 'w').write(asmcfg.dot()) print "generating IR... %x" % start_addr ircfg = ir_arch.new_ircfg_from_asmcfg(asmcfg) if verbose: print "IR ok... %x" % start_addr for irb in ircfg.blocks.itervalues(): irs = [] for assignblk in irb: new_assignblk = { expr_simp(dst): expr_simp(src) for dst, src in assignblk.iteritems() } irs.append(AssignBlock(new_assignblk, instr=assignblk.instr)) ircfg.blocks[irb.loc_key] = IRBlock(irb.loc_key, irs) if verbose: out = ircfg.dot() open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out) title = "Miasm IR graph" if simplify: dead_simp(ir_arch, ircfg) ircfg.simplify(expr_simp) modified = True while modified: modified = False modified |= dead_simp(ir_arch, ircfg) modified |= remove_empty_assignblks(ircfg) modified |= merge_blocks(ircfg, entry_points) title += " (simplified)" if type_graph == TYPE_GRAPH_IR: graph = GraphMiasmIR(ircfg, title, None) graph.Show() return head = list(entry_points)[0] class IRAOutRegs(ira): def get_out_regs(self, block): regs_todo = super(IRAOutRegs, self).get_out_regs(block) out = {} for assignblk in block: for dst in assignblk: reg = self.ssa_var.get(dst, None) if reg is None: continue if reg in regs_todo: out[reg] = dst return set(out.values()) # Add dummy dependency to uncover out regs affectation for loc in ircfg.leaves(): irblock = ircfg.blocks.get(loc) if irblock is None: continue regs = {} for reg in ir_arch.get_out_regs(irblock): regs[reg] = reg assignblks = list(irblock) new_assiblk = AssignBlock(regs, assignblks[-1].instr) assignblks.append(new_assiblk) new_irblock = IRBlock(irblock.loc_key, assignblks) ircfg.blocks[loc] = new_irblock ir_arch = IRAOutRegs(mdis.loc_db) ir_arch.ssa_var = {} modified = True ssa_forbidden_regs = set( [ir_arch.pc, ir_arch.IRDst, ir_arch.arch.regs.exception_flags]) head = list(entry_points)[0] heads = set([head]) all_ssa_vars = {} propagate_expr = PropagateExpr() ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.ssa_variable_to_expr.update(all_ssa_vars) ssa.transform(head) all_ssa_vars.update(ssa.ssa_variable_to_expr) ir_arch.ssa_var.update(ssa.ssa_variable_to_expr) if simplify: while modified: ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.ssa_variable_to_expr.update(all_ssa_vars) ssa.transform(head) all_ssa_vars.update(ssa.ssa_variable_to_expr) ir_arch.ssa_var.update(ssa.ssa_variable_to_expr) while modified: modified = False modified |= propagate_expr.propagate(ssa, head) modified |= ircfg.simplify(expr_simp) simp_modified = True while simp_modified: simp_modified = False simp_modified |= dead_simp(ir_arch, ircfg) simp_modified |= remove_empty_assignblks(ircfg) simp_modified |= load_from_int(ircfg, bs, is_addr_ro_variable) modified |= simp_modified ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.ssa_variable_to_expr.update(all_ssa_vars) ssa.transform(head) all_ssa_vars.update(ssa.ssa_variable_to_expr) if type_graph == TYPE_GRAPH_IRSSA: graph = GraphMiasmIR(ssa.graph, title, None) graph.Show() return if type_graph == TYPE_GRAPH_IRSSAUNSSA: cfg_liveness = DiGraphLivenessSSA(ssa.graph) cfg_liveness.init_var_info(ir_arch) cfg_liveness.compute_liveness() UnSSADiGraph(ssa, head, cfg_liveness) if simplify: modified = True while modified: modified = False modified |= ssa.graph.simplify(expr_simp) simp_modified = True while simp_modified: simp_modified = False simp_modified |= dead_simp(ir_arch, ssa.graph) simp_modified |= remove_empty_assignblks(ssa.graph) simp_modified |= merge_blocks(ssa.graph, heads) modified |= simp_modified graph = GraphMiasmIR(ssa.graph, title, None) graph.Show()
if args.verbose > 3: open('tmp_before_%d.dot' % index, 'w').write(ircfg_a.dot()) simp_modified = False log.info("dead simp...") simp_modified |= dead_simp(ir_arch_a, ircfg_a) log.info("ok...") index += 1 if args.verbose > 3: open('tmp_after_%d.dot' % index, 'w').write(ircfg_a.dot()) simp_modified |= remove_empty_assignblks(ircfg_a) simp_modified |= del_unused_edges(ircfg_a, heads) simp_modified |= merge_blocks(ircfg_a, heads) if args.loadint: simp_modified |= load_from_int(ircfg_a, bs, is_addr_ro_variable) modified |= simp_modified index += 1 if args.verbose > 3: open('stack_%d.dot' % index, 'w').write(ircfg_a.dot()) if args.stack2var: modified |= replace_stack_vars(ir_arch_a, ssa) if args.verbose > 3: open('final_pre.dot', 'w').write(ircfg_a.dot()) if args.verbose > 3: open('final_merge.dot', 'w').write(ircfg_a.dot()) ssa = SSADiGraph(ircfg_a) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.ssa_variable_to_expr.update(all_ssa_vars)
def build_graph(verbose=False, simplify=False, ssa=False, ssa_simplify=False): start_addr = idc.ScreenEA() machine = guess_machine(addr=start_addr) dis_engine, ira = 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.loc_db) # populate symbols with ida names for addr, name in idautils.Names(): if name is None: continue if (mdis.loc_db.get_offset_location(addr) or mdis.loc_db.get_name_location(name)): # Symbol alias continue mdis.loc_db.add_location(name, addr) if verbose: print "start disasm" if verbose: print hex(addr) asmcfg = mdis.dis_multiblock(start_addr) entry_points = set([mdis.loc_db.get_offset_location(start_addr)]) if verbose: print "generating graph" open('asm_flow.dot', 'w').write(asmcfg.dot()) print "generating IR... %x" % start_addr ircfg = ir_arch.new_ircfg_from_asmcfg(asmcfg) if verbose: print "IR ok... %x" % start_addr for irb in ircfg.blocks.itervalues(): irs = [] for assignblk in irb: new_assignblk = { expr_simp(dst): expr_simp(src) for dst, src in assignblk.iteritems() } irs.append(AssignBlock(new_assignblk, instr=assignblk.instr)) ircfg.blocks[irb.loc_key] = IRBlock(irb.loc_key, irs) if verbose: out = ircfg.dot() open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out) title = "Miasm IR graph" if simplify: dead_simp(ir_arch, ircfg) ircfg.simplify(expr_simp) modified = True while modified: modified = False modified |= dead_simp(ir_arch, ircfg) modified |= remove_empty_assignblks(ircfg) modified |= merge_blocks(ircfg, entry_points) title += " (simplified)" graph = GraphMiasmIR(ircfg, title, None) if ssa: if len(entry_points) != 1: raise RuntimeError("Your graph should have only one head") head = list(entry_points)[0] ssa = SSADiGraph(ircfg) ssa.transform(head) title += " (SSA)" graph = GraphMiasmIR(ssa.graph, title, None) if ssa_simplify: class IRAOutRegs(ira): def get_out_regs(self, block): regs_todo = super(self.__class__, self).get_out_regs(block) out = {} for assignblk in block: for dst in assignblk: reg = self.ssa_var.get(dst, None) if reg is None: continue if reg in regs_todo: out[reg] = dst return set(out.values()) # Add dummy dependecy to uncover out regs affectation for loc in ircfg.leaves(): irblock = ircfg.blocks.get(loc) if irblock is None: continue regs = {} for reg in ir_arch.get_out_regs(irblock): regs[reg] = reg assignblks = list(irblock) new_assiblk = AssignBlock(regs, assignblks[-1].instr) assignblks.append(new_assiblk) new_irblock = IRBlock(irblock.loc_key, assignblks) ircfg.blocks[loc] = new_irblock ir_arch = IRAOutRegs(mdis.loc_db) def is_addr_ro_variable(bs, addr, size): """ Return True if address at @addr is a read-only variable. WARNING: Quick & Dirty @addr: integer representing the address of the variable @size: size in bits """ try: _ = bs.getbytes(addr, size/8) except IOError: return False return True ir_arch.ssa_var = {} index = 0 modified = True ssa_forbidden_regs = set([ ir_arch.pc, ir_arch.IRDst, ir_arch.arch.regs.exception_flags ]) head = list(entry_points)[0] heads = set([head]) all_ssa_vars = set() propagate_expr = PropagateExpr() while modified: ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.transform(head) all_ssa_vars.update(ssa._ssa_variable_to_expr) ssa_regs = [reg for reg in ssa.expressions if reg.is_id()] ssa_forbidden_regs.update(ssa_regs) ir_arch.ssa_var.update(ssa._ssa_variable_to_expr) while modified: index += 1 modified = False modified |= propagate_expr.propagate(ssa, head) modified |= ircfg.simplify(expr_simp) simp_modified = True while simp_modified: index += 1 simp_modified = False simp_modified |= dead_simp(ir_arch, ircfg) index += 1 simp_modified |= remove_empty_assignblks(ircfg) simp_modified |= merge_blocks(ircfg, heads) simp_modified |= load_from_int(ircfg, bs, is_addr_ro_variable) modified |= simp_modified index += 1 merge_blocks(ircfg, heads) ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.transform(head) all_ssa_vars.update(ssa._ssa_variable_to_expr) ssa._ssa_variable_to_expr = all_ssa_vars dead_simp(ir_arch, ssa.graph) title += " (SSA Simplified)" graph = GraphMiasmIR(ssa.graph, title, None) graph.Show()
def build_graph(verbose=False, simplify=False, ssa=False, ssa_simplify=False): start_addr = idc.ScreenEA() machine = guess_machine(addr=start_addr) dis_engine, ira = 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.loc_db) # populate symbols with ida names for addr, name in idautils.Names(): if name is None: continue if (mdis.loc_db.get_offset_location(addr) or mdis.loc_db.get_name_location(name)): # Symbol alias continue mdis.loc_db.add_location(name, addr) if verbose: print "start disasm" if verbose: print hex(addr) asmcfg = mdis.dis_multiblock(start_addr) entry_points = set([mdis.loc_db.get_offset_location(start_addr)]) if verbose: print "generating graph" open('asm_flow.dot', 'w').write(asmcfg.dot()) print "generating IR... %x" % start_addr ircfg = ir_arch.new_ircfg_from_asmcfg(asmcfg) if verbose: print "IR ok... %x" % start_addr for irb in ircfg.blocks.itervalues(): irs = [] for assignblk in irb: new_assignblk = { expr_simp(dst): expr_simp(src) for dst, src in assignblk.iteritems() } irs.append(AssignBlock(new_assignblk, instr=assignblk.instr)) ircfg.blocks[irb.loc_key] = IRBlock(irb.loc_key, irs) if verbose: out = ircfg.dot() open(os.path.join(tempfile.gettempdir(), 'graph.dot'), 'wb').write(out) title = "Miasm IR graph" if simplify: dead_simp(ir_arch, ircfg) ircfg.simplify(expr_simp) modified = True while modified: modified = False modified |= dead_simp(ir_arch, ircfg) modified |= remove_empty_assignblks(ircfg) modified |= merge_blocks(ircfg, entry_points) title += " (simplified)" graph = GraphMiasmIR(ircfg, title, None) if ssa: if len(entry_points) != 1: raise RuntimeError("Your graph should have only one head") head = list(entry_points)[0] ssa = SSADiGraph(ircfg) ssa.transform(head) title += " (SSA)" graph = GraphMiasmIR(ssa.graph, title, None) if ssa_simplify: class IRAOutRegs(ira): def get_out_regs(self, block): regs_todo = super(self.__class__, self).get_out_regs(block) out = {} for assignblk in block: for dst in assignblk: reg = self.ssa_var.get(dst, None) if reg is None: continue if reg in regs_todo: out[reg] = dst return set(out.values()) # Add dummy dependency to uncover out regs affectation for loc in ircfg.leaves(): irblock = ircfg.blocks.get(loc) if irblock is None: continue regs = {} for reg in ir_arch.get_out_regs(irblock): regs[reg] = reg assignblks = list(irblock) new_assiblk = AssignBlock(regs, assignblks[-1].instr) assignblks.append(new_assiblk) new_irblock = IRBlock(irblock.loc_key, assignblks) ircfg.blocks[loc] = new_irblock ir_arch = IRAOutRegs(mdis.loc_db) def is_addr_ro_variable(bs, addr, size): """ Return True if address at @addr is a read-only variable. WARNING: Quick & Dirty @addr: integer representing the address of the variable @size: size in bits """ try: _ = bs.getbytes(addr, size / 8) except IOError: return False return True ir_arch.ssa_var = {} index = 0 modified = True ssa_forbidden_regs = set( [ir_arch.pc, ir_arch.IRDst, ir_arch.arch.regs.exception_flags]) head = list(entry_points)[0] heads = set([head]) all_ssa_vars = set() propagate_expr = PropagateExpr() while modified: ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.transform(head) all_ssa_vars.update(ssa._ssa_variable_to_expr) ssa_regs = [reg for reg in ssa.expressions if reg.is_id()] ssa_forbidden_regs.update(ssa_regs) ir_arch.ssa_var.update(ssa._ssa_variable_to_expr) while modified: index += 1 modified = False modified |= propagate_expr.propagate(ssa, head) modified |= ircfg.simplify(expr_simp) simp_modified = True while simp_modified: index += 1 simp_modified = False simp_modified |= dead_simp(ir_arch, ircfg) index += 1 simp_modified |= remove_empty_assignblks(ircfg) simp_modified |= merge_blocks(ircfg, heads) simp_modified |= load_from_int(ircfg, bs, is_addr_ro_variable) modified |= simp_modified index += 1 merge_blocks(ircfg, heads) ssa = SSADiGraph(ircfg) ssa.immutable_ids.update(ssa_forbidden_regs) ssa.transform(head) all_ssa_vars.update(ssa._ssa_variable_to_expr) ssa._ssa_variable_to_expr = all_ssa_vars dead_simp(ir_arch, ssa.graph) title += " (SSA Simplified)" graph = GraphMiasmIR(ssa.graph, title, None) graph.Show()
modified |= ircfg_a.simplify(expr_simp) if args.verbose > 3: open('tmp_simp_%d.dot' % index, 'w').write(ircfg_a.dot()) simp_modified = True while simp_modified: index += 1 if args.verbose > 3: open('tmp_before_%d.dot' % index, 'w').write(ircfg_a.dot()) simp_modified = False simp_modified |= dead_simp(ir_arch_a, ircfg_a) index += 1 if args.verbose > 3: open('tmp_after_%d.dot' % index, 'w').write(ircfg_a.dot()) simp_modified |= remove_empty_assignblks(ircfg_a) simp_modified |= merge_blocks(ircfg_a, heads) simp_modified |= load_from_int(ircfg_a, bs, is_addr_ro_variable) modified |= simp_modified index += 1 if args.verbose > 3: open('stack_%d.dot' % index, 'w').write(ircfg_a.dot()) if args.stack2var: modified |= replace_stack_vars(ir_arch_a, ssa) if args.verbose > 3: open('final_pre.dot', 'w').write(ircfg_a.dot()) merge_blocks(ircfg_a, heads) if args.verbose > 3: open('final_merge.dot', 'w').write(ircfg_a.dot()) ssa = SSADiGraph(ircfg_a) ssa.immutable_ids.update(ssa_forbidden_regs)