from idaapi import Form
from idaapi import GraphViewer
from idaapi import plugin_t
import re

def getText(ea):
    t=GetFunctionName(ea)
    if t=='':
        t='%s:%x: %s'%(SegName(ea),ea,GetDisasm(ea))
    return t

class Xrefs:
    def __init__(self,ea,max_depth=1,regexp=None):
        self.xrefs_to={}
        self.xrefs_from={}
        self.ea=ea
        self.max_depth=max_depth
        self.regexp=regexp

    def _recGetXrefsTo(self,ea,depth=1):
        found=False
        self.xrefs_to[ea]=[]
        for xref in XrefsTo(ea,idaapi.XREF_FAR):
            if not xref.iscode:
                continue
            x=xref.frm
            f=idaapi.get_func(x)
            if f:
                x=f.startEA
            if x in self.xrefs_to[ea]:
                continue
            t=getText(x)
            if self.regexp!=None:
                if re.search(self.regexp,t)!=None:
                    self.xrefs_to[ea].append(x)
                    found=True
                    continue
                if f and depth<self.max_depth \
                    and x not in self.xrefs_to \
                    and self._recGetXrefsTo(x,depth+1)==True:
                    self.xrefs_to[ea].append(x)
                    found=True
                elif x in self.xrefs_to and len(self.xrefs_to[x])>0:
                    self.xrefs_to[ea].append(x)
                    found=True
            else:
                self.xrefs_to[ea].append(x)
                if f and depth<self.max_depth and x not in self.xrefs_to:
                    self._recGetXrefsTo(x,depth+1)
        return found

    def getXrefsTo(self):
        self.xrefs_to={}
        f=idaapi.get_func(self.ea)
        if f:
            self._recGetXrefsTo(f.startEA)
        return self.xrefs_to

    def _recGetXrefsFrom(self,ea,depth=1):
        found=False
        self.xrefs_from[ea]=[]
        for x in [x for x in FuncItems(ea) if idaapi.is_call_insn(x)]:
            for xref in XrefsFrom(x,idaapi.XREF_FAR):
                if not xref.iscode:
                    continue
                if xref.to in self.xrefs_from[ea]:
                    continue
                t=getText(xref.to)
                if self.regexp!=None:
                    if re.search(self.regexp,t)!=None:
                        self.xrefs_from[ea].append(xref.to)
                        found=True
                        continue
                    if depth<self.max_depth and xref.to not in self.xrefs_from \
                        and self._recGetXrefsFrom(xref.to,depth+1)==True:
                        self.xrefs_from[ea].append(xref.to)
                        found=True
                    elif xref.to in self.xrefs_from and \
                        len(self.xrefs_from[xref.to])>0:
                        self.xrefs_from[ea].append(xref.to)
                        found=True
                else:
                    self.xrefs_from[ea].append(xref.to)
                    if depth<self.max_depth and xref.to not in self.xrefs_from:
                        self._recGetXrefsFrom(xref.to,depth+1)
        return found

    def getXrefsFrom(self):
        self.xrefs_from={}
        f=idaapi.get_func(self.ea)
        if f:
            self._recGetXrefsFrom(f.startEA)
        return self.xrefs_from

class XrefsGraph(GraphViewer):
    def __init__(self,title,ea,xrefs_to,xrefs_from):
        GraphViewer.__init__(self,title)
        self.xrefs_to=xrefs_to
        self.xrefs_from=xrefs_from

    def OnRefresh(self):
        self.Clear()
        nodes={}
        for ea in self.xrefs_to:
            if len(self.xrefs_to[ea])==0:
                continue
            if ea not in nodes:
                nodes[ea]=self.AddNode(ea)
            for x in self.xrefs_to[ea]:
                if x not in nodes:
                    nodes[x]=self.AddNode(x)
                self.AddEdge(nodes[x],nodes[ea])
        for ea in self.xrefs_from:
            if len(self.xrefs_from[ea])==0:
                continue
            if ea not in nodes:
                nodes[ea]=self.AddNode(ea)
            for x in self.xrefs_from[ea]:
                if x not in nodes:
                    nodes[x]=self.AddNode(x)
                self.AddEdge(nodes[ea],nodes[x])
        return True

    def OnGetText(self,node_id):
        return getText(self[node_id])

    def OnDblClick(self,node_id):
        ea=self[node_id]
        idc.Jump(ea)
        return True

class XrefsForm(Form):
    def __init__(self):
        Form.__init__(self,r"""STARTITEM {id:iAddr}
BUTTON YES* Ok
BUTTON CANCEL Cancel
Xrefs Graph

{FormChangeCb}
<#Address must be within a function#Enter an address:{iAddr}>

Directions:
<Cross references from:{rFrom}>
<Cross references to:{rTo}>{cGroupDirections}>

Options:
<Recursion:{rRecursion}>
<Include data references (NOT USED YET):{rDrefs}>
<Apply regular expression filter:{rRegexp}>
<Display recursion dots (NOT USED YET):{rDots}>
<Display comments (NOT USED YET):{rComments}>{cGroupOptions}>

<##Recursion depth:{iDepth}>
<##Regular expression:{iRegexp}>
""", {
    'iAddr': Form.NumericInput(tp=Form.FT_ADDR),
    'cGroupDirections': Form.ChkGroupControl(('rFrom','rTo')),
    'cGroupOptions': Form.ChkGroupControl(('rRecursion',
                                           'rDrefs',
                                           'rRegexp',
                                           'rDots',
                                           'rComments')),
    'iDepth': Form.NumericInput(tp=Form.FT_UINT64),
    'iRegexp': Form.StringInput(),
    'FormChangeCb': Form.FormChangeCb(self.OnFormChange)
    })

    def OnFormChange(self,fid):
        print('>>fid=%d'%(fid))
        return 1

class XrefsGraphPlugin_t(plugin_t):
    flags=idaapi.PLUGIN_UNL
    comment="Xrefs Graph plugin for IDA"
    help=""
    wanted_name="Xrefs Graph"
    wanted_hotkey="Alt-F8"

    def init(self):
        return idaapi.PLUGIN_OK

    def run(self,arg=0):
        f=XrefsForm()
        f.Compile()
        f.iAddr.value=ScreenEA()
        f.iDepth.value=16
        f.cGroupDirections.value=3
        f.cGroupOptions.value=5
        f.iRegexp.value='^xxx'
        ok=f.Execute()
        #print('>>ok=%d'%(ok))
        if ok==1:
            #print('f.cGroupDirections.value=%x'%(f.cGroupDirections.value))
            #print('f.cGroupOptions.value=%x'%(f.cGroupOptions.value))
            ea=f.iAddr.value
            max_depth=f.iDepth.value
            if f.cGroupOptions.value&1==0:
                max_depth=1
            regexp=None
            if f.cGroupOptions.value&4!=0:
                regexp=f.iRegexp.value
            x=Xrefs(ea,max_depth,regexp)
            xrefs_to={}
            xrefs_from={}
            if f.cGroupDirections.value&2!=0:
                xrefs_to.update(x.getXrefsTo())
            if f.cGroupDirections.value&1!=0:
                xrefs_from.update(x.getXrefsFrom())
            title='XrefsGraph for %s (depth=%d,regexp=%s)'%(getText(ea),
                                                            max_depth,regexp)
            g=XrefsGraph(title,ea,xrefs_to,xrefs_from)
            g.Show()
        f.Free()
        return

    def term(self):
        return

def PLUGIN_ENTRY():
    return XrefsGraphPlugin_t()