121 lines
3.8 KiB
Python
121 lines
3.8 KiB
Python
import idaapi
|
|
import idc
|
|
|
|
|
|
def inverse_if_condition(cif):
|
|
# cexpr_t has become broken but fortunately still exist `assing` method which copies one expr into another
|
|
cit_if_condition = cif.expr
|
|
tmp_cexpr = idaapi.cexpr_t()
|
|
tmp_cexpr.assign(cit_if_condition)
|
|
new_if_condition = idaapi.lnot(tmp_cexpr)
|
|
cif.expr.swap(new_if_condition)
|
|
del cit_if_condition
|
|
|
|
|
|
def inverse_if(cif):
|
|
inverse_if_condition(cif)
|
|
idaapi.qswap(cif.ithen, cif.ielse)
|
|
|
|
|
|
class InversionInfo(object):
|
|
ARRAY_NAME = "$HexRaysPyTools:IfThenElse:"
|
|
|
|
def __init__(self, func_ea):
|
|
self.__name = InversionInfo.ARRAY_NAME + hex(int(func_ea))
|
|
self.__id = idc.GetArrayId(self.__name)
|
|
|
|
def get_inverted(self):
|
|
if self.__id != -1:
|
|
array = idc.GetArrayElement(idc.AR_STR, self.__id, 0)
|
|
return set(map(int, array.split()))
|
|
return set()
|
|
|
|
def switch_inverted(self, address):
|
|
if self.__id == -1:
|
|
self.__id = idc.CreateArray(self.__name)
|
|
idc.SetArrayString(self.__id, 0, str(address))
|
|
else:
|
|
inverted = self.get_inverted()
|
|
try:
|
|
inverted.remove(address)
|
|
if not inverted:
|
|
idc.DeleteArray(self.__id)
|
|
|
|
except KeyError:
|
|
inverted.add(address)
|
|
|
|
idc.SetArrayString(self.__id, 0, " ".join(map(str, inverted)))
|
|
|
|
|
|
class SpaghettiVisitor(idaapi.ctree_parentee_t):
|
|
def __init__(self):
|
|
super(SpaghettiVisitor, self).__init__()
|
|
|
|
def visit_insn(self, instruction):
|
|
if instruction.op != idaapi.cit_block:
|
|
return 0
|
|
|
|
while True:
|
|
cblock = instruction.cblock
|
|
size = cblock.size()
|
|
# Find block that has "If" and "return" as last 2 statements
|
|
if size < 2:
|
|
break
|
|
|
|
if cblock.at(size - 2).op != idaapi.cit_if:
|
|
break
|
|
|
|
cif = cblock.at(size - 2).cif
|
|
if cblock.back().op != idaapi.cit_return or cif.ielse:
|
|
break
|
|
|
|
cit_then = cif.ithen
|
|
|
|
# Skip if only one (not "if") statement in "then" branch
|
|
if cit_then.cblock.size() == 1 and cit_then.cblock.front().op != idaapi.cit_if:
|
|
return 0
|
|
|
|
inverse_if_condition(cif)
|
|
|
|
# Take return from list of statements and later put it back
|
|
cit_return = idaapi.cinsn_t()
|
|
cit_return.assign(instruction.cblock.back())
|
|
cit_return.thisown = False
|
|
instruction.cblock.pop_back()
|
|
|
|
# Fill main block with statements from "Then" branch
|
|
while cit_then.cblock:
|
|
instruction.cblock.push_back(cit_then.cblock.front())
|
|
cit_then.cblock.pop_front()
|
|
|
|
# Put back main return if there's no another return or "GOTO" already
|
|
if instruction.cblock.back().op not in (idaapi.cit_return, idaapi.cit_goto):
|
|
new_return = idaapi.cinsn_t()
|
|
new_return.thisown = False
|
|
new_return.assign(cit_return)
|
|
instruction.cblock.push_back(new_return)
|
|
|
|
# Put return into "Then" branch
|
|
cit_then.cblock.push_back(cit_return)
|
|
return 0
|
|
|
|
|
|
class SwapThenElseVisitor(idaapi.ctree_parentee_t):
|
|
def __init__(self, func_ea):
|
|
super(SwapThenElseVisitor, self).__init__()
|
|
self.__inversion_info = InversionInfo(func_ea)
|
|
self.__inverted = self.__inversion_info.get_inverted()
|
|
|
|
def visit_insn(self, insn):
|
|
if insn.op != idaapi.cit_if or insn.cif.ielse is None:
|
|
return 0
|
|
|
|
if insn.ea in self.__inverted:
|
|
inverse_if(insn.cif)
|
|
|
|
return 0
|
|
|
|
def apply_to(self, *args):
|
|
if self.__inverted:
|
|
super(SwapThenElseVisitor, self).apply_to(*args)
|