Xrefs Graph

from idaapi import Form
from idaapi import GraphViewer
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
            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
                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* Yes
BUTTON NO Nope
BUTTON CANCEL Nevermind
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

def main():
    f=XrefsForm()
    f.Compile()
    f.iAddr.value=here()
    f.iDepth.value=4
    f.cGroupDirections.value=3
    f.cGroupOptions.value=3
    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()

main()