""" Invert the then and else blocks of a cif_t.

Author: EiNSTeiN_ <einstein@g3nius.org>

This is a rewrite in Python of the vds3 example that comes with hexrays sdk.


The main difference with the original C code is that when we create the inverted
condition object, the newly created cexpr_t instance is given to the hexrays and
must not be freed by swig. To achieve this, we have to change the 'thisown' flag
when appropriate. See http://www.swig.org/Doc1.3/Python.html#Python_nn35

"""

import idautils
import idaapi
import idc

import traceback

NETNODE_NAME = '$ hexrays-inverted-if'

class hexrays_callback_info(object):

    def __init__(self):
        self.vu = None

        self.node = idaapi.netnode()
        if not self.node.create(NETNODE_NAME):
            # node exists
            self.load()
        else:
            self.stored = []

        return

    def load(self):

        self.stored = []

        try:
            data = self.node.getblob(0, 'I')
            if data:
                self.stored = eval(data)
                print 'Invert-if: Loaded %s' % (repr(self.stored), )
        except:
            print 'Failed to load invert-if locations'
            traceback.print_exc()
            return

        return

    def save(self):

        try:
            self.node.setblob(repr(self.stored), 0, 'I')
        except:
            print 'Failed to save invert-if locations'
            traceback.print_exc()
            return

        return

    def invert_if(self, cfunc, insn):

        if insn.opname != 'if':
            return False

        cif = insn.details

        if not cif.ithen or not cif.ielse:
            return False

        idaapi.qswap(cif.ithen, cif.ielse)
        cond = idaapi.cexpr_t(cif.expr)
        notcond = idaapi.lnot(cond)
        cond.thisown = 0 # the new wrapper 'notcond' now holds the reference to the cexpr_t

        cif.expr.swap(notcond)

        return True

    def add_location(self, ea):
        if ea in self.stored:
            self.stored.remove(ea)
        else:
            self.stored.append(ea)
        self.save()
        return

    def find_if_statement(self, vu):

        vu.get_current_item(idaapi.USE_KEYBOARD)
        item = vu.item

        if item.is_citem() and item.it.op == idaapi.cit_if and item.it.to_specific_type.cif.ielse is not None:
            return item.it.to_specific_type

        if vu.tail.citype == idaapi.VDI_TAIL and vu.tail.loc.itp == idaapi.ITP_ELSE:
            # for tail marks, we know only the corresponding ea,
            # not the pointer to if-statement
            # find it by walking the whole ctree
            class if_finder_t(idaapi.ctree_visitor_t):
                def __init__(self, ea):
                    idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST | idaapi.CV_INSNS)

                    self.ea = ea
                    self.found = None
                    return

                def visit_insn(self, i):
                    if i.op == idaapi.cit_if and i.ea == self.ea:
                        self.found = i
                        return 1 # stop enumeration
                    return 0

            iff = if_finder_t(vu.tail.loc.ea)
            if iff.apply_to(vu.cfunc.body, None):
                return iff.found

        return

    def invert_if_event(self, vu):

        cfunc = vu.cfunc.__deref__()

        i = self.find_if_statement(vu)
        if not i:
            return False

        if self.invert_if(cfunc, i):
            vu.refresh_ctext()

            self.add_location(i.ea)

        return True

    def restore(self, cfunc):

        class visitor(idaapi.ctree_visitor_t):

            def __init__(self, inverter, cfunc):
                idaapi.ctree_visitor_t.__init__(self, idaapi.CV_FAST | idaapi.CV_INSNS)
                self.inverter = inverter
                self.cfunc = cfunc
                return

            def visit_insn(self, i):
                try:
                    if i.op == idaapi.cit_if and i.ea in self.inverter.stored:
                        self.inverter.invert_if(self.cfunc, i)
                except:
                    traceback.print_exc()
                return 0 # continue enumeration

        visitor(self, cfunc).apply_to(cfunc.body, None)

        return

    def menu_callback(self):
        try:
            self.invert_if_event(self.vu)
        except:
            traceback.print_exc()
        return 0

    def event_callback(self, event, *args):

        try:
            if event == idaapi.hxe_keyboard:
                vu, keycode, shift = args

                if idaapi.lookup_key_code(keycode, shift, True) == idaapi.get_key_code("I") and shift == 0:
                    if self.invert_if_event(vu):
                        return 1

            elif event == idaapi.hxe_right_click:
                self.vu,  = args
                idaapi.add_custom_viewer_popup_item(self.vu.ct, "Invert then/else", "I", self.menu_callback)

            elif event == idaapi.hxe_maturity:
                cfunc, maturity = args

                if maturity == idaapi.CMAT_FINAL:
                    self.restore(cfunc)
        except:
            traceback.print_exc()

        return 0

if idaapi.init_hexrays_plugin():
    i = hexrays_callback_info()
    idaapi.install_hexrays_callback(i.event_callback)
else:
    print 'invert-if: hexrays is not available.'