def reconstruct_while_loop(self): # remove the branch going into the loop and leave only a way to exit the loop. if len(self.loop.start.container) == 1: block = self.loop.condition_block stmt = block.container[-1] condition = stmt.expr dest, = list(set(block.jump_to).difference(self.loop.exits)) stmt.remove() block.container.add(goto_t(stmt.ea, value_t(dest.ea, self.function.arch.address_size))) else: condition = value_t(1, 1) block = self.loop.condition_block stmt = block.container[-1] dest, = list(set(block.jump_to).difference(self.loop.exits)) stmt.remove() ctn = container_t(block, [break_t(stmt.ea)]) _if = if_t(stmt.ea, b_not_t(stmt.expr.copy()), ctn) block.container.add(_if) simplify_expressions.run(_if.expr, deep=True) # collapse all the loop blocks into a single container #print repr(self.loop.blocks) self.cf.reconstruct_forward(self.loop.blocks, self.prioritize_longest_path) #print repr(self.loop.blocks) # build the new while loop. block = self.loop.blocks[0] _while = self.wrap_loop(stmt.ea, while_t, block, condition.copy()) block.container.add(goto_t(None, value_t(self.loop.exit_block.ea, self.function.arch.address_size))) self.cleanup_loop(_while, block, self.loop.exit_block) return
def convert_if_branch(flow, block, container): """ very simple if() form. """ for stmt in container: if type(stmt) != branch_t: continue true_block = flow.blocks[stmt.true.value] false_block = flow.blocks[stmt.false.value] if type(true_block.container[-1]) == goto_t and \ true_block.container[-1].expr.value == stmt.false.value and \ len(true_block.jump_from) == 1: newblock = if_t(stmt.expr.pluck(), container_t(true_block.container[:-1])) simplify_expressions.run(newblock.expr, deep=True) block.container.insert(stmt.index(), newblock) block.container.insert(stmt.index(), goto_t(stmt.false)) stmt.remove() false_block.jump_from.remove(true_block) block.jump_to.remove(flow.blocks.pop(stmt.true.value)) return True if type(false_block.container[-1]) == goto_t and \ false_block.container[-1].expr.value == stmt.true.value and \ len(false_block.jump_from) == 1: newblock = if_t(b_not_t(stmt.expr.pluck()), container_t(false_block.container[:-1])) simplify_expressions.run(newblock.expr, deep=True) block.container.insert(stmt.index(), newblock) block.container.insert(stmt.index(), goto_t(stmt.true)) stmt.remove() true_block.jump_from.remove(false_block) block.jump_to.remove(flow.blocks.pop(stmt.false.value)) return True return False
def switch_goto_if_needed(block, dstblock): """ if the last item at the end of 'block' is a goto to dstblock, do nothing, otherwise invert that goto with the one in the if_t in the next-to-last position. """ container = block.container assert type(container[-1]) == goto_t if container[-1].expr.value == dstblock.ea: return if len(container) < 2: return assert type(container[-2]) == if_t assert len(container[-2].then_expr) == 1 assert type(container[-2].then_expr[0]) == goto_t assert container[-2].then_expr[0].expr.value == dstblock.ea # invert goto_t destinations container[-1].expr.value, container[-2].then_expr[0].expr.value = \ container[-2].then_expr[0].expr.value, container[-1].expr.value container[-2].expr = b_not_t(container[-2].expr.copy()) simplify_expressions.run(container[-2].expr, deep=True) return
def invert_goto_condition(stmt): """ invert the goto at the end of a block for the goto in the if_t preceding it """ stmt.true.value, stmt.false.value = stmt.false.value, stmt.true.value stmt.expr = b_not_t(stmt.expr.pluck()) simplify_expressions.run(stmt.expr, deep=True) return
def invert_goto_condition(block): """ invert the goto at the end of a block for the goto in the if_t preceding it """ stmt = block.container[-2] stmt.then_expr[0], block.container[-1] = block.container[-1], stmt.then_expr[0] stmt.expr = b_not_t(stmt.expr.copy()) simplify_expressions.run(stmt.expr, deep=True) return
def expand_branches(self, blocks=None): for stmt in iterators.statement_iterator_t(self.function): if type(stmt) != branch_t: continue if blocks and stmt.container.block not in blocks: continue condition = stmt.expr.copy() goto_true = goto_t(stmt.ea, stmt.true.copy()) goto_false = goto_t(stmt.ea, stmt.false.copy()) _if = if_t(stmt.ea, condition, container_t(stmt.container.block, [goto_true])) simplify_expressions.run(_if.expr, deep=True) stmt.container.insert(stmt.index(), _if) stmt.container.insert(stmt.index(), goto_false) stmt.remove() return
def run(self): self.cf.reconstruct_forward(self.conditional.left, self.prioritizer) self.cf.reconstruct_forward(self.conditional.right, self.prioritizer) if len(self.conditional.left) == 0 and len(self.conditional.right) == 0: return if type(self.conditional.top.container[-1]) not in (goto_t, branch_t): return if self.conditional.top in [loop.start for loop in self.cf.loops]: return if len(self.conditional.left) > 0: then_blocks, else_blocks = self.conditional.left, self.conditional.right else: then_blocks, else_blocks = self.conditional.right, self.conditional.left expr = self.conditional_expr(self.conditional.top, then_blocks[0]) stmt = self.conditional.top.container[-1] stmt.remove() prioritized = False if self.prioritizer and len(then_blocks) > 0 and len(else_blocks) > 0: first = self.prioritizer(then_blocks[0], else_blocks[0]) if first is else_blocks[0]: then_blocks, else_blocks = else_blocks, then_blocks expr = b_not_t(expr) prioritized = True then_ctn = self.cf.assembler.build_container(container_t(self.conditional.top), then_blocks, self.prioritizer) else_ctn = self.cf.assembler.build_container(container_t(self.conditional.top), else_blocks, self.prioritizer) if not prioritized and else_ctn and type(then_ctn[0]) in (if_t, branch_t) and type(else_ctn[0]) not in (if_t, branch_t): then_ctn, else_ctn = else_ctn, then_ctn expr = b_not_t(expr) _if = if_t(stmt.ea, expr, then_ctn, else_ctn) simplify_expressions.run(_if.expr, deep=True) self.conditional.top.container.add(_if) self.conditional.top.container.add(goto_t(stmt.ea, value_t(self.conditional.bottom.ea, self.function.arch.address_size))) self.remove_goto(then_ctn, self.conditional.bottom) if else_ctn: self.remove_goto(else_ctn, self.conditional.bottom) return
def combine_branch_blocks(flow, this, next): """ combine two if_t that jump to the same destination into a boolean or expression. """ left = [this.container[-1].true.value, this.container[-1].false.value] right = [next.container[-1].true.value, next.container[-1].false.value] dest = list(set(left).intersection(set(right))) if len(dest) == 1: # both blocks have one jump in common. dest = dest[0] if this.container[-1].false.value == dest: invert_goto_condition(this.container[-1]) if next.container[-1].false.value == dest: invert_goto_condition(next.container[-1]) common = flow.blocks[dest] exit = flow.blocks[next.container[-1].false.value] if exit == this: cls = b_and_t else: cls = b_or_t stmt = this.container[-1] stmt.expr = cls(stmt.expr.copy(), next.container[-1].expr.copy()) simplify_expressions.run(stmt.expr, deep=True) this.container[-1].false = next.container[-1].false this.jump_to.remove(next) next.jump_from.remove(this) flow.blocks[dest].jump_from.remove(next) exit.jump_from.remove(next) if exit != this: exit.jump_from.append(this) this.jump_to.append(exit) return True return False
def combine_ifs(flow, block, container): """ process if_t """ for stmt in container: # invert then and else side if then-side is empty if type(stmt) == if_t and stmt.else_expr is not None and len(stmt.then_expr) == 0: stmt.then_expr = stmt.else_expr stmt.expr = b_not_t(stmt.expr.copy()) stmt.else_expr = None simplify_expressions.run(stmt.expr, deep=True) return True # remove if altogether if it contains no statements at all if type(stmt) == if_t and stmt.else_expr is None and len(stmt.then_expr) == 0: container.remove(stmt) return True return False
def convert_elseif(flow, block, container): """ if we have an if_t as only statement in the then-side of a parent if_t, and the parent if_t has an else-side which doesn't contain an if_t as only statement (to avoid infinite loops), then we can safely invert the two sides of the parent if_t so that it will be displayed in the more natural 'if(...) { } else if(...) {}' form. """ for stmt in container: if type(stmt) == if_t and stmt.else_expr and \ len(stmt.then_expr) == 1 and type(stmt.then_expr[0]) == if_t and \ not (len(stmt.else_expr) == 1 and type(stmt.else_expr[0]) == if_t): \ stmt.then_expr, stmt.else_expr = stmt.else_expr, stmt.then_expr stmt.expr = b_not_t(stmt.expr.copy()) simplify_expressions.run(stmt.expr, deep=True) return True return False
def remove_goto(self, ctn, block): """ remove goto going to block at the end of the given container """ stmt = ctn[-1] if type(stmt) == goto_t and stmt.expr.value == block.ea: stmt.remove() elif type(stmt) == branch_t: if stmt.true.value == block.ea: condition = b_not_t(stmt.expr.pluck()) goto = goto_t(None, stmt.false.copy()) elif stmt.false.value == block.ea: condition = stmt.expr.pluck() goto = goto_t(None, stmt.true.copy()) else: return _if = if_t(stmt.ea, condition, container_t(block, [goto])) simplify_expressions.run(_if.expr, deep=True) ctn.add(_if) stmt.remove() self.connect_next(_if.then_expr, []) self.remove_goto(_if.then_expr, block) return
def reconstruct_while_loop(self): # remove the branch going into the loop and leave only a way to exit the loop. if len(self.loop.start.container) == 1: block = self.loop.condition_block stmt = block.container[-1] condition = stmt.expr dest, = list(set(block.jump_to).difference(self.loop.exits)) stmt.remove() block.container.add( goto_t(stmt.ea, value_t(dest.ea, self.function.arch.address_size))) else: condition = value_t(1, 1) block = self.loop.condition_block stmt = block.container[-1] dest, = list(set(block.jump_to).difference(self.loop.exits)) stmt.remove() ctn = container_t(block, [break_t(stmt.ea)]) _if = if_t(stmt.ea, b_not_t(stmt.expr.copy()), ctn) block.container.add(_if) simplify_expressions.run(_if.expr, deep=True) # collapse all the loop blocks into a single container #print repr(self.loop.blocks) self.cf.reconstruct_forward(self.loop.blocks, self.prioritize_longest_path) #print repr(self.loop.blocks) # build the new while loop. block = self.loop.blocks[0] _while = self.wrap_loop(stmt.ea, while_t, block, condition.copy()) block.container.add( goto_t( None, value_t(self.loop.exit_block.ea, self.function.arch.address_size))) self.cleanup_loop(_while, block, self.loop.exit_block) return
def combine_ifs(flow, block, container): """ process if_t """ for stmt in container: # invert then and else side if then-side is empty if type(stmt) == if_t and stmt.else_expr is not None and len( stmt.then_expr) == 0: stmt.then_expr = stmt.else_expr stmt.expr = b_not_t(stmt.expr.copy()) stmt.else_expr = None simplify_expressions.run(stmt.expr, deep=True) return True # remove if altogether if it contains no statements at all if type(stmt) == if_t and stmt.else_expr is None and len( stmt.then_expr) == 0: container.remove(stmt) return True return False
def connect_next(self, container, blocks, prioritizer=None, exclude=[]): stmt = container[-1] if type(stmt) == goto_t and stmt.expr.value in self.function.blocks: dest = self.function.blocks[stmt.expr.value] if dest in blocks or (len(list(dest.jump_from)) == 1 and not dest.node.is_return_node and dest not in exclude): stmt.remove() self.assemble_connected(container, blocks, dest, prioritizer, exclude) elif type(stmt) == branch_t: dest_true, dest_false = None, None if stmt.true.value in self.function.blocks: dest_true = self.function.blocks[stmt.true.value] if stmt.false.value in self.function.blocks: dest_false = self.function.blocks[stmt.false.value] expr = stmt.expr.copy() if prioritizer and dest_true and dest_false: first = prioritizer(dest_true, dest_false) if first is dest_false: dest_true, dest_false = dest_false, dest_true expr = b_not_t(expr) true_ctn = container_t(container.block, []) _if = if_t(stmt.ea, expr, true_ctn, None) simplify_expressions.run(_if.expr, deep=True) container.add(_if) stmt.remove() if dest_true and (dest_true in blocks or (len(list(dest_true.jump_from)) == 1 and not dest_true.node.is_return_node and dest_true not in exclude)): self.assemble_connected(true_ctn, blocks, dest_true, prioritizer, exclude) else: true_ctn.add(goto_t(None, stmt.true.copy())) if dest_false and (dest_false in blocks or (len(list(dest_false.jump_from)) == 1 and not dest_false.node.is_return_node and dest_true not in exclude)): self.assemble_connected(container, blocks, dest_false, prioritizer, exclude) else: container.add(goto_t(None, stmt.false.copy())) return
def combine_branch_blocks(cls, function, this, next): """ combine two if_t that jump to the same destination into a boolean or expression. """ left = [this.container[-1].true.value, this.container[-1].false.value] right = [next.container[-1].true.value, next.container[-1].false.value] dest = list(set(left).intersection(set(right))) if len(dest) != 1: return False # both blocks have one jump in common. dest = dest[0] if this.container[-1].false.value == dest: cls.invert_goto_condition(this.container[-1]) if next.container[-1].false.value == dest: cls.invert_goto_condition(next.container[-1]) common = function.blocks[dest] exit = function.blocks[next.container[-1].false.value] if exit == this: cls = b_and_t else: cls = b_or_t stmt = this.container[-1] stmt.expr = cls(stmt.expr.copy(), next.container[-1].expr.copy()) simplify_expressions.run(stmt.expr, deep=True) this.container[-1].false = next.container[-1].false function.blocks.pop(next.ea) return True
def combine_if_blocks(flow, this, next): """ combine two if_t that jump to the same destination into a boolean or expression. """ left = [this.container[-1].expr.value, this.container[-2].then_expr[0].expr.value] right = [next.container[-1].expr.value, next.container[-2].then_expr[0].expr.value] dest = list(set(left).intersection(set(right))) if len(dest) == 1: # both blocks have one jump in common. dest = dest[0] if this.container[-1].expr.value == dest: invert_goto_condition(this) if next.container[-1].expr.value == dest: invert_goto_condition(next) other = flow.blocks[next.container[-1].expr.value] if other == this: cls = b_and_t else: cls = b_or_t stmt = this.container[-2] expr = cls(stmt.expr, next.container[-2].expr) stmt.expr = simplify_expressions.run(expr, deep=True) this.jump_to.remove(next) next.jump_from.remove(this) flow.blocks[dest].jump_from.remove(next) other.jump_from.remove(next) if other != this: other.jump_from.append(this) this.jump_to.append(other) this.container[-1] = next.container[-1] return True return False