def fix_multiple_next_constraints(asmcfg, mode): """ When there are multiple blocks proceeding another block with no jump, add one. :return: """ blocks_to_be_added = [] for loc_key in asmcfg.nodes(): next_edges = { edge: constraint for edge, constraint in asmcfg.edges2constraint.items() if constraint == AsmConstraint.c_next } pred_next = list(ploc_key for (ploc_key, dloc_key) in next_edges if dloc_key == loc_key) if len(pred_next) > 1: for index in range(1, len(pred_next)): inst = create_jump_instruction(mode, ExprLoc(loc_key, mode)) new_block_loc_key = asmcfg.loc_db.add_location() new_block = AsmBlock(new_block_loc_key) new_block.addline(inst) new_block.bto = {AsmConstraintTo(loc_key)} asmcfg.loc_key_to_block(pred_next[index]).bto = { AsmConstraintNext(new_block_loc_key) } blocks_to_be_added.append(new_block) # one while might be sufficient, depends on type of _nodes for block in blocks_to_be_added: asmcfg.add_block(block)
def _insert_flat_block(self, source_flat_block, symb_exec, flat_block_to_loc_key): """ Copies source_flat_block and sets its successors according to flat_block_to_loc_key :param flat_block_to_loc_key: dictionary mapping flat_blocks to respective loc_keys :param symb_exec: instance of current symbolic execution engine :param source_flat_block: flat_block to be inserted :return: dictionary mapping old successor loc_keys to the new ones """ # we're not using redirect_successors after copying to avoid executing the same loops multiple times source_block = self.asmcfg.loc_key_to_block( source_flat_block.block_loc_key) tobe_processed = {} new_flat_blocks = set() new_block_loc_key = flat_block_to_loc_key[source_flat_block] if self.out_asmcfg.loc_key_to_block(new_block_loc_key) is not None: raise Exception("Target loc_key is already associated to a block") new_block = AsmBlock(new_block_loc_key) # copy instructions for ln in source_block.lines: tmp_ln = instruction_x86(ln.name, ln.mode, [i.copy() for i in ln.args], ln.additional_info) tmp_ln.b = ln.b tmp_ln.l = ln.l tmp_ln.offset = ln.offset new_block.addline(tmp_ln) constraints = source_block.bto # try to simplify the destination if it's a primary flattening block if not self.flat_loops[source_block.loc_key].is_default: logger.debug("current block is a part of primary loc_keys") simplified_target = symb_exec.eval_expr(self.ircfg.IRDst) if isinstance(simplified_target, ExprInt): simplified_target = self.asmcfg.loc_db.get_offset_location( int(simplified_target)) elif isinstance(simplified_target, ExprLoc): simplified_target = simplified_target.loc_key else: # there's probably a(n) (series of) unknown instruction(s) causing an implicit conditional assignment # such as CMOV or SBB->AND->ADD, prepend comparison + cond jump if it happens to be common, or add it to # ExtendedAsmCFG.extended_discovery and split flow on the final instruction # it's also possible that it's not related to any cff loop at all addr = self.asmcfg.loc_db.get_location_offset( source_flat_block.block_loc_key) addr = hex(addr) if addr else addr logger.warning( "Couldn't simplify loc_key %s at %s, continuing" % (str(source_flat_block.block_loc_key), addr)) logger.warning("the simplified target is %s of instance %s" % (simplified_target, type(simplified_target))) simplified_target = None if simplified_target: constraints = {AsmConstraintTo(simplified_target)} mode = self.asmcfg.mode # remove redundant comparison dp = DependencyGraph(self.ircfg, True) block_loc_key = source_block.loc_key res = next( dp.get(block_loc_key, {self.ircfg.IRDst}, None, {block_loc_key})) for depnode in res.relevant_nodes: ind = depnode.line_nb ind -= (len(self.ircfg.blocks[block_loc_key]) - len(new_block.lines)) if new_block.lines[ind].name == "CMP": new_block.lines.pop(ind) new_block.lines[-1] = create_jump_instruction( mode, ExprLoc(simplified_target, mode)) # copy constraints new_bto = set() for constraint in constraints: if not self.asmcfg.loc_key_to_block(constraint.loc_key): logger.debug("Skipping bad constraint %s" % constraint.loc_key) continue flat_block = self.flat_loops.get_block(constraint.loc_key, symb_exec, source_flat_block) if flat_block not in flat_block_to_loc_key: new_flat_blocks.add(flat_block) new_loc_key = self.out_asmcfg.loc_db.add_location() tobe_processed[constraint.loc_key] = (new_loc_key, flat_block) flat_block_to_loc_key[flat_block] = new_loc_key else: new_loc_key = flat_block_to_loc_key[flat_block] new_bto.add(AsmConstraint(new_loc_key, constraint.c_t)) new_block.bto = new_bto new_block.alignment = source_block.alignment # change jmp targets if new_block.lines: for ind, arg in enumerate(list(new_block.lines[-1].args)): if isinstance(arg, ExprLoc): if not self.asmcfg.loc_key_to_block(arg.loc_key): logger.debug("Skipping bad constraint %s" % arg.loc_key) continue new_target, flat_block = tobe_processed.get( arg.loc_key, (None, None)) if not new_target: flat_block = self.flat_loops.get_block( arg.loc_key, symb_exec, source_flat_block) new_target = flat_block_to_loc_key.get(flat_block) # None in case of irrelevant calls logger.debug("new target: %s" % new_target) if new_target: new_block.lines[-1].args[ind] = ExprLoc( new_target, arg.size) self.out_asmcfg.add_block(new_block) return new_flat_blocks