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_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 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 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
