""" Xref plugin for Hexrays Decompiler

Author: EiNSTeiN_ <einstein@g3nius.org>

Show decompiler-style Xref when the X key is pressed in the Decompiler window.

- It supports any global name: functions, strings, integers, etc.
- It supports structure member.

"""

import idautils
import idaapi
import idc

import traceback

try:
    from PyQt4 import QtCore, QtGui
    print 'Using PyQt'
except:
    print 'PyQt not available'
    
    try:
        from PySide import QtGui, QtCore
        print 'Using PySide'
    except:
        print 'PySide not available'

XREF_EA = 0
XREF_STRUC_MEMBER = 1

class XrefsForm(idaapi.PluginForm):
    
    def __init__(self, target):
        
        idaapi.PluginForm.__init__(self)
        
        self.target = target
        
        if type(self.target) == idaapi.cfunc_t:
            
            self.__type = XREF_EA
            self.__ea = self.target.entry_ea
            self.__name = 'Xrefs of %x' % (self.__ea, )
            
        elif type(self.target) == idaapi.cexpr_t and self.target.opname == 'obj':
            
            self.__type = XREF_EA
            self.__ea = self.target.obj_ea
            self.__name = 'Xrefs of %x' % (self.__ea, )
        
        elif type(self.target) == idaapi.cexpr_t and self.target.opname in ('memptr', 'memref'):
            
            self.__type = XREF_STRUC_MEMBER
            name = self.get_struc_name()
            self.__name = 'Xrefs of %s' % (name, )
        
        else:
            raise ValueError('cannot show xrefs for this kind of target')
        
        return
    
    def get_struc_name(self):
        
        x = self.target.operands['x']
        m = self.target.operands['m']
        
        xtype = typestring(x.type.u_str())
        xtype.remove_ptr_or_array()
        typename = str(xtype)
        
        sid = idc.GetStrucIdByName(typename)
        member = idc.GetMemberName(sid, m)
        
        return '%s::%s' % (typename, member)
    
    def OnCreate(self, form):
        
        # Get parent widget
        try:
            self.parent = self.FormToPySideWidget(form)
        except:
            self.parent = self.FormToPyQtWidget(form)
        
        self.populate_form()
        
        return
    
    def Show(self):
        idaapi.PluginForm.Show(self, self.__name)
        return
    
    def populate_form(self):
        # Create layout
        layout = QtGui.QVBoxLayout()

        layout.addWidget(QtGui.QLabel(self.__name))
        self.table = QtGui.QTableWidget()
        layout.addWidget(self.table)
        
        self.table.setColumnCount(3)
        self.table.setHorizontalHeaderItem(0, QtGui.QTableWidgetItem("Address"))
        self.table.setHorizontalHeaderItem(1, QtGui.QTableWidgetItem("Function"))
        self.table.setHorizontalHeaderItem(2, QtGui.QTableWidgetItem("Line"))
        
        self.table.setColumnWidth(0, 80)
        self.table.setColumnWidth(1, 150)
        self.table.setColumnWidth(2, 450)
        
        self.table.cellDoubleClicked.connect(self.double_clicked)
        
        #~ self.table.setSelectionMode(QtGui.QAbstractItemView.NoSelection)
        self.table.setSelectionBehavior(QtGui.QAbstractItemView.SelectRows )
        self.parent.setLayout(layout)
        
        self.populate_table()
        
        return
    
    def double_clicked(self, row, column):
        
        ea = self.functions[row]
        idaapi.open_pseudocode(ea, True)
        
        return
    
    def get_decompiled_line(self, cfunc, ea):
        
        print repr(ea)
        if ea not in cfunc.eamap:
            print 'strange, %x is not in %x eamap' % (ea, cfunc.entry_ea)
            return
        
        insnvec = cfunc.eamap[ea]
        
        lines = []
        for stmt in insnvec:
            
            qs = idaapi.qstring()
            qp = idaapi.qstring_printer_t(cfunc.__deref__(), qs, False)
            
            stmt._print(0, qp)
            s = str(qs).split('\n')[0]
                
            #~ s = idaapi.tag_remove(s)
            lines.append(s)
        
        return '\n'.join(lines)
    
    def get_items_for_ea(self, ea):
        
        frm = [x.frm for x in idautils.XrefsTo(self.__ea)]
        
        items = []
        for ea in frm:
            try:
                cfunc = idaapi.decompile(ea)
                cfunc.refcnt += 1
                
                self.functions.append(cfunc.entry_ea)
                self.items.append((ea, idc.GetFunctionName(cfunc.entry_ea), self.get_decompiled_line(cfunc, ea)))
                
            except Exception as e:
                print 'could not decompile: %s' % (str(e), )
                raise
        
        return
    
    def get_items_for_type(self):
        
        x = self.target.operands['x']
        m = self.target.operands['m']
        
        xtype = typestring(x.type.u_str())
        xtype.remove_ptr_or_array()
        typename = str(xtype)
        
        addresses = []
        for ea in idautils.Functions():
            
            try:
                cfunc = idaapi.decompile(ea)
                cfunc.refcnt += 1
            except:
                print 'Decompilation of %x failed' % (ea, )
                continue
            
            str(cfunc)
            
            for citem in cfunc.treeitems:
                citem = citem.to_specific_type
                if not (type(citem) == idaapi.cexpr_t and citem.opname in ('memptr', 'memref')):
                    continue
                
                _x = citem.operands['x']
                _m = citem.operands['m']
                _xtype = typestring(_x.type.u_str())
                _xtype.remove_ptr_or_array()
                _typename = str(_xtype)
                
                #~ print 'in', hex(cfunc.entry_ea), _typename, _m
                
                if not (_typename == typename and _m == m):
                    continue
                
                parent = citem
                while parent:
                    if type(parent.to_specific_type) == idaapi.cinsn_t:
                        break
                    parent = cfunc.body.find_parent_of(parent)
                
                if not parent:
                    print 'cannot find parent statement (?!)'
                    continue
                
                if parent.ea in addresses:
                    continue
                
                if parent.ea == idaapi.BADADDR:
                    print 'parent.ea is BADADDR'
                    continue
                
                addresses.append(parent.ea)
                
                self.functions.append(cfunc.entry_ea)
                self.items.append((parent.ea, idc.GetFunctionName(cfunc.entry_ea), self.get_decompiled_line(cfunc, int(parent.ea))))
            
        
        return []
    
    def populate_table(self):
        
        self.functions = []
        self.items = []
        
        if self.__type == XREF_EA:
            self.get_items_for_ea(self.__ea)
        else:
            self.get_items_for_type()
        
        self.table.setRowCount(len(self.items))
        
        i = 0
        for item in self.items:
            address, func, line = item
            item = QtGui.QTableWidgetItem('0x%x' % (address, ))
            item.setFlags(item.flags() ^ QtCore.Qt.ItemIsEditable)
            self.table.setItem(i, 0, item)
            item = QtGui.QTableWidgetItem(func)
            item.setFlags(item.flags() ^ QtCore.Qt.ItemIsEditable)
            self.table.setItem(i, 1, item)
            item = QtGui.QTableWidgetItem(line)
            item.setFlags(item.flags() ^ QtCore.Qt.ItemIsEditable)
            self.table.setItem(i, 2, item)
            
            i += 1
        
        self.table.resizeRowsToContents()
        
        return
    
    def OnClose(self, form):
        pass

class hexrays_callback_info(object):
    
    def __init__(self):
        self.vu = None
        return
    
    def show_xrefs(self, vu):
        
        vu.get_current_item(idaapi.USE_KEYBOARD)
        item = vu.item
        
        sel = None
        if item.citype == idaapi.VDI_EXPR and item.it.to_specific_type.opname in ('obj', 'memref', 'memptr'):
            # if an expression is selected. verify that it's either a cot_obj, cot_memref or cot_memptr
            sel = item.it.to_specific_type
            
        elif item.citype == idaapi.VDI_FUNC:
            # if the function itself is selected, show xrefs to it.
            sel = item.f
        else:
            return False
        
        form = XrefsForm(sel)
        form.Show()
        return True
    
    def menu_callback(self):
        self.show_xrefs(self.vu)
        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("X") and shift == 0:
                    if self.show_xrefs(vu):
                        return 1
                
            elif event == idaapi.hxe_right_click:
                self.vu = args[0]
                idaapi.add_custom_viewer_popup_item(self.vu.ct, "Xrefs", "X", self.menu_callback)
            
        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.'