[
  {
    "path": ".gitignore",
    "content": "*.pyproj\n*.sln\n.vs/*\nidaapi.py\n.env\n.vscode\n*.pyc\n.DS_Store\n"
  },
  {
    "path": "LICENSE.txt",
    "content": "This software is provided 'as-is', without any express or implied\nwarranty. In no event will the authors be held liable for any damages\narising from the use of this software.\n\nPermission is granted to anyone to use this software for any purpose,\nincluding commercial applications, and to alter it and redistribute it\nfreely, subject to the following restrictions:\n\n   1. The origin of this software must not be misrepresented; you must not\n   claim that you wrote the original software. If you use this software\n   in a product, an acknowledgment in the product documentation would be\n   appreciated but is not required.\n\n   2. Altered source versions must be plainly marked as such, and must not be\n   misrepresented as being the original software.\n\n   3. This notice may not be removed or altered from any source\n   distribution.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\nANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\nWARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY\nDIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\nLOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\nON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\nSOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
  },
  {
    "path": "classdiagram.py",
    "content": "from idaapi import GraphViewer, ask_file\n\n# The  below will only be displayed as bases\nignore_namespaces = (\"std\", \"type_info\")\n\nclass ClassDiagram(GraphViewer):\n\n    def __init__(self, title, classes):\n        self.classes = self.transitive_reduction(classes)\n        GraphViewer.__init__(self, title)\n\n    def transitive_reduction(self, graph):\n        for u in graph.keys():\n            print('u node: ' + u + '(parents: ' + ', '.join([v for v in graph[u]]) + ')')\n            for v in graph[u]:\n                # Compute the dfs from v\n                print('DFS from v node: ' + v)\n                dfs = self.dfs_paths(graph, v)\n                for node in dfs:\n                    print('  - DFS node: ' + str(node))\n                    if v != node and node in graph[u]:\n                        graph[u].remove(node)\n                        print('  # Removed ' + u + ' -> ' + node)\n        return graph\n\n    def dfs_paths(self, graph, start, path=None):\n        if path is None:\n            path = [start]\n        # check for leaf\n        if start not in graph.keys() or len(graph[start])==0:\n            #print('    + ' + str(len(graph[start])))\n            #print(start + ': []')\n            for p in path:\n                yield p\n        else:\n            #print(start + ': [' + ', '.join([s for s in graph[start]]) + ']')\n            print('    + RECUR: ' + ', '.join(set(graph[start]) - set(path)))\n            for next in set(graph[start]) - set(path):\n                for p in self.dfs_paths(graph, next, path + [next]):\n                    print('    + ' + p)\n                    yield p\n\n    def add_node(self, class_name):\n        if class_name is None:\n            return\n        if not class_name.startswith(ignore_namespaces):\n            print(\"Adding node: \" + class_name)\n            new_node = self.AddNode(class_name)\n            self.name_to_node[class_name] = new_node\n\n    def OnRefresh(self):\n        self.Clear()\n        self.name_to_node = {}\n        # Create nodes\n        for class_name in self.classes.keys():\n            # Skipping common namespaces\n            self.add_node(class_name)\n        # Create edges\n        for class_name in self.name_to_node.keys():\n            print(\"Adding edges for: \" + class_name)\n            node = self.name_to_node[class_name]\n            print(\"bases: \" + str(self.classes[class_name]))\n            for base_name in self.classes[class_name]:\n                if base_name not in self.name_to_node:\n                    # Add originally skipped base node\n                    self.add_node(base_name)\n                base = self.name_to_node.get(base_name)\n                if base is not None:\n                    self.AddEdge(base, node)\n        return True\n\n    def OnGetText(self, node_id):\n          return self[node_id]\n\n    # dot file export modified from http://joxeankoret.com\n    def OnCommand(self, cmd_id):\n        if self.cmd_dot == cmd_id:\n            fname = ask_file(1, \"*.dot\", \"Export DOT file\")\n            if fname:\n                f = open(fname, \"wt\")\n                buf = \"digraph G {\\n graph [overlap=scale]; node [fontname=Courier]; rankdir=\\\"LR\\\";\\n\\n\"\n                keyList = list(self.classes.keys())\n                for c in keyList:\n                    n = keyList.index(c)\n                    buf += ' a%s [shape=box, label = \"%s\", color=\"blue\"]\\n' % (n, c)\n                buf += \"\\n\"\n                for c in keyList:\n                    class_index = keyList.index(c)\n                    for base in self.classes[c]:\n                        if base in keyList:\n                            base_index = keyList.index(base)\n                            buf += ' a%s -> a%s [style = bold]\\n' % (class_index, base_index)\n                buf += \"}\"\n                f.write(buf)\n                f.close()\n\n    def Show(self):\n      if not GraphViewer.Show(self):\n          return False\n      self.cmd_dot = self.AddCommand(\"Export DOT\", \"F2\")\n      return True"
  },
  {
    "path": "classinformer.py",
    "content": "# ClassInformer python\n# Nicolas Guigo / NCC Group\n# Tyler Colgan / NCC Group\n# 03/2017\n\nimport idaapi\nfrom idc import *\nfrom ida_search import find_text\n\nidaapi.require(\"utils\")\nidaapi.require(\"msvc\")\nidaapi.require(\"gcc\")\nidaapi.require(\"classdiagram\")\nfrom idaapi import auto_is_ok\nfrom msvc import run_msvc\nfrom gcc import run_gcc\nfrom classdiagram import ClassDiagram\n\ndef show_classes(classes):\n    c = ClassDiagram(\"Class Diagram\", classes)\n    c.Show()\n\ndef isGcc():\n    gcc_info = find_text(0x0, 0, 0, \"N10__cxxabiv117__class_type_infoE\", SEARCH_CASE|SEARCH_DOWN)\n    return gcc_info != BADADDR\n\ndef main():\n    print(\"Starting ClassInformerPython\")\n    if auto_is_ok():\n        classes = run_gcc() if isGcc() else run_msvc()\n        print(classes)\n        show_classes(classes)\n    else:\n        print(\"Take it easy, man\")\n    print(\"Done\")\n\nif __name__ == '__main__':\n    main()\n"
  },
  {
    "path": "gcc.py",
    "content": "# Modified from GCC RTTI parsing code originally written by Igor Skochinsky.\n# See http://www.hexblog.com/?p=704 for the original version of his code.\n\nimport idaapi\nfrom idaapi import BADADDR\nfrom idc import *\nfrom utils import utils\nu = utils()\n\nall_classes = {}\n\nti_names = [\n \"St9type_info\",\n \"N10__cxxabiv117__class_type_infoE\",\n \"N10__cxxabiv120__si_class_type_infoE\",\n \"N10__cxxabiv121__vmi_class_type_infoE\",\n]\n\nTI_TINFO = 0\nTI_CTINFO = 1\nTI_SICTINFO = 2\nTI_VMICTINFO = 3\n\nclass BaseClass:\n    def __init__(self, ti, offset, flags):\n        self.ti = ti\n        self.offset = offset\n        self.flags = flags\n\nclass ClassDescriptor:\n    def __init__(self, vtable, namestr):\n        self.vtable = vtable\n        self.namestr = namestr\n        self.bases = []\n\n    def add_base(self, base, offset=0, flags=0):\n        self.bases.append(BaseClass(base, offset, flags))\n\ndef tinfo2class(tiname):\n  s = demangle_name(tiname, 0)\n  if s is None:\n      return s\n  return s.replace(\"`typeinfo for'\",\"\")\n\ndef classname(namestr):\n  return tinfo2class(\"__ZTI\" + namestr)\n\n# dd `vtable for'std::type_info+8\n# dd `typeinfo name for'std::type_info\ndef format_type_info(ea):\n  # get the class name string\n  tis = u.get_ptr(ea + u.PTR_SIZE)\n  if u.is_bad_addr(tis):\n    return BADADDR\n  bname = get_strlit_contents(tis)\n  if bname == None or len(bname) == 0:\n    return BADADDR\n  # looks good, let's do it\n  name = bname.decode('UTF-8')\n  ea2 = u.format_struct(ea, \"vp\")\n  u.force_name(tis, \"__ZTS\" + name)\n  u.force_name(ea, \"__ZTI\" + name)\n  # find our vtable\n  # 0 followed by ea\n  pat = u.ptr_to_bytes(0) + \" \" + u.ptr_to_bytes(ea)\n  vtb = find_binary(0, SEARCH_CASE|SEARCH_DOWN, pat)\n  if not u.is_bad_addr(vtb):\n    print(\"vtable for %s at %08X\" % (name, vtb))\n    u.format_struct(vtb, \"pp\")\n    u.force_name(vtb, u.vtname(name))\n  else:\n    vtb = BADADDR\n  all_classes[ea] = ClassDescriptor(vtb, name)\n  return ea2\n\n# dd `vtable for'__cxxabiv1::__si_class_type_info+8\n# dd `typeinfo name for'MyClass\n# dd `typeinfo for'BaseClass\ndef format_si_type_info(ea):\n  ea2 = format_type_info(ea)\n  pbase = u.get_ptr(ea2)\n  all_classes[ea].add_base(pbase)\n  ea2 = u.format_struct(ea2, \"p\")\n  return ea2\n\n# dd `vtable for'__cxxabiv1::__si_class_type_info+8\n# dd `typeinfo name for'MyClass\n# dd flags\n# dd base_count\n# (base_type, offset_flags) x base_count\ndef format_vmi_type_info(ea):\n  ea2 = format_type_info(ea)\n  ea2 = u.format_struct(ea2, \"ii\")\n  base_count = create_data(ea2-4, FF_DWORD, 4, ida_idaapi.BADADDR)\n  clas = all_classes[ea]\n  if base_count > 100:\n    print(\"%08X: over 100 base classes?!\" % ea)\n    return BADADDR\n  for i in range(base_count):\n    base_ti = u.get_ptr(ea2)\n    flags_off = u.get_ptr(ea2 + u.PTR_SIZE)\n    off = u.SIGNEXT(flags_off>>8, 24)\n    clas.add_base(base_ti, off, flags_off & 0xFF)\n    ea2 = u.format_struct(ea2, \"pl\")\n  return ea2\n\ndef find_type_info(idx):\n  name = ti_names[idx]\n  ea = u.find_string(name)\n  if ea != BADADDR:\n    xrefs = u.xref_or_find(ea)\n    if xrefs:\n      ti_start = xrefs[0] - u.PTR_SIZE\n      if not u.is_bad_addr(ti_start):\n        print(\"found %d at %08X\" % (idx, ti_start))\n        ea2 = format_type_info(ti_start)\n        if idx >= TI_CTINFO:\n          u.format_struct(ea2, \"p\")\n\ndef handle_classes(idx, formatter):\n  name = u.vtname(ti_names[idx])\n  ea = get_name_ea_simple(name)\n  if ea == BADADDR:\n    # try single underscore\n    name = name[1:]\n    ea = get_name_ea_simple(name)\n  if ea == BADADDR:\n    print(\"Could not find vtable for %s\" % ti_names[idx])\n    return\n  idx = 0\n  handled = set()\n  while ea != BADADDR:\n    print(\"Looking for refs to vtable %08X\" % ea)\n    if idaapi.is_spec_ea(ea):\n      xrefs = u.xref_or_find(ea, True)\n      ea += u.PTR_SIZE*2\n      xrefs.extend(u.xref_or_find(ea, True))\n    else:\n      ea += u.PTR_SIZE*2\n      xrefs = u.xref_or_find(ea, True)\n    for x in xrefs:\n      if not u.is_bad_addr(x) and not x in handled:\n        print(\"found %s at %08X\" % (name, x))\n        ea2 = formatter(x)\n        handled.add(x)\n    ea = get_name_ea_simple(\"%s_%d\" % (name, idx))\n    idx += 1\n\ndef run_gcc():\n    classes = {}\n    # turn on GCC3 demangling\n    idaapi.cvar.inf.demnames |= idaapi.DEMNAM_GCC3\n    print(\"Looking for standard type info classes\")\n    find_type_info(TI_TINFO)\n    find_type_info(TI_CTINFO)\n    find_type_info(TI_SICTINFO)\n    find_type_info(TI_VMICTINFO)\n    print(\"Looking for simple classes\")\n    handle_classes(TI_CTINFO, format_type_info)\n    print(\"Looking for single-inheritance classes\")\n    handle_classes(TI_SICTINFO, format_si_type_info)\n    print(\"Looking for multiple-inheritance classes\")\n    handle_classes(TI_VMICTINFO, format_vmi_type_info)\n    for i in range(len(all_classes)):\n        tiaddr = u.num2key(all_classes)[i]\n        klass = all_classes[tiaddr]\n        name = classname(klass.namestr)\n        ti = \"%08X\" % tiaddr\n        vt = \"%08X\" % klass.vtable\n        basestr = []\n        for b in klass.bases:\n            if b.ti in all_classes:\n                bklass = all_classes[b.ti]\n                basename = classname(bklass.namestr)\n            elif idaapi.is_spec_ea(b.ti):\n                nm = get_name(b.ti, ida_name.GN_VISIBLE)\n                basename = tinfo2class(nm)\n            else:\n                print(\"Base %08X not found for class %08X!\" % (b.ti, tiaddr))\n                basename = \"ti_%08X\" % b.ti\n            basestr.append(basename)\n        classes[name] = basestr\n        print(\"basestr: \\\"%s\\\"\" % basestr)\n    return classes"
  },
  {
    "path": "msvc.py",
    "content": "from idaapi import get_struc_id, BADADDR, del_struc, get_struc, add_struc, add_struc_member, FF_DATA, FF_DWORD, FF_0OFF, get_struc_size, FF_STRLIT, del_items, DELIT_DELNAMES, create_struct, get_member_by_name, get_32bit, get_strlit_contents, demangle_name, create_dword, op_offset\nfrom idc import *\nfrom utils import utils\nu = utils()\n\nclasses = {}\n\nclass RTTIStruc:\n    tid = 0\n    struc = 0\n    size = 0\n\ndef strip(name):\n    if name.startswith(\"class \") and name.endswith(\"`RTTI Type Descriptor'\"):\n        return name[6:-23]\n    elif name.startswith(\"struct \") and name.endswith(\"`RTTI Type Descriptor'\"):\n        return name[7:-23]\n    else:\n        return name\n\nclass RTTICompleteObjectLocator(RTTIStruc):\n\n    # Init class statics\n    msid = get_struc_id(\"RTTICompleteObjectLocator\")\n    if msid != BADADDR:\n        del_struc(msid)\n    msid = add_struc(0xFFFFFFFF, \"RTTICompleteObjectLocator\", False)\n    add_struc_member(msid, \"signature\", BADADDR, FF_DATA|FF_DWORD, -1, 4)\n    add_struc_member(msid, \"offset\", BADADDR, FF_DATA|FF_DWORD, -1, 4)\n    add_struc_member(msid, \"cdOffset\", BADADDR, FF_DATA|FF_DWORD, -1, 4)\n    add_struc_member(msid, \"pTypeDescriptor\", BADADDR, FF_DATA|FF_DWORD|FF_0OFF, u.mt_rva().tid, 4)\n    add_struc_member(msid, \"pClassDescriptor\", BADADDR, FF_DATA|FF_DWORD|FF_0OFF, u.mt_rva().tid, 4)\n    if u.x64:\n        add_struc_member(msid, \"pSelf\", BADADDR, FF_DATA|FF_DWORD|FF_0OFF, u.mt_rva().tid, 4)\n    tid = msid\n    struc = get_struc(tid)\n    size = get_struc_size(tid)\n    print(\"Completed Registering RTTICompleteObjectLocator\")\n\n    def __init__(self, ea, vtable):\n        del_items(ea, DELIT_DELNAMES, self.size)\n        if ida_bytes.create_struct(ea, self.size, self.tid):\n            # Get adress of type descriptor from CompleteLocator\n            print(\"Complete Object Locator at: 0x%x\" % ea)\n            offset = get_member_by_name(self.struc, \"pTypeDescriptor\").soff\n            typeDescriptor = get_32bit(ea+offset) + u.x64_imagebase()\n            print(\"Looking for type Descriptor at: 0x%x\" % typeDescriptor)\n            rtd = RTTITypeDescriptor(typeDescriptor)\n            if rtd.class_name:\n                print(\"Type Descriptor at: 0x%x\" % typeDescriptor)\n                offset = get_member_by_name(self.struc, \"pClassDescriptor\").soff\n                classHierarchyDes = get_32bit(ea+offset) + u.x64_imagebase()\n                rchd = RTTIClassHierarchyDescriptor(classHierarchyDes)\n                # filter out None entries\n                rchd.bases = filter(lambda x: x, rchd.bases)\n                classes[strip(rtd.class_name)] = [strip(b) for b in rchd.bases]\n                set_name(vtable, \"vtable__\" + strip(rtd.class_name), SN_NOWARN)\n            else:\n                # if the RTTITypeDescriptor doesn't have a valid name for us to\n                # read, then this wasn't a valid RTTICompleteObjectLocator\n                del_items(ea, self.size, DELIT_SIMPLE)\n\nclass RTTITypeDescriptor(RTTIStruc):\n    class_name = None\n\n    msid = get_struc_id(\"RTTITypeDescriptor\")\n    if msid != BADADDR:\n        del_struc(msid)\n    msid = add_struc(0xFFFFFFFF, \"RTTITypeDescriptor\", False)\n    add_struc_member(msid, \"pVFTable\", BADADDR, FF_DATA|u.PTR_TYPE|FF_0OFF, u.mt_address().tid, u.PTR_SIZE)\n    add_struc_member(msid, \"spare\", BADADDR, FF_DATA|u.PTR_TYPE, -1, u.PTR_SIZE)\n    add_struc_member(msid, \"name\", BADADDR, FF_DATA|FF_STRLIT, u.mt_ascii().tid, 0)\n    tid = msid\n    struc = get_struc(tid)\n    size = get_struc_size(tid)\n    print(\"Completed Registering RTTITypeDescriptor\")\n\n    def __init__(self, ea):\n        name = ea + get_member_by_name(get_struc(self.tid), \"name\").soff\n        strlen = u.get_strlen(name)\n        if strlen is None:\n            # not a real vtable\n            return\n        self.size = self.size + strlen\n        bmangled = get_strlit_contents(name, strlen, 0)\n        if bmangled is None:\n            # not a real function name\n            return\n        mangled = bmangled.decode('UTF-8')\n        print(\"Mangled: \" + mangled)\n        demangled = demangle_name('??_R0' + mangled[1:] , 0)\n        if demangled:\n            del_items(ea, DELIT_DELNAMES, self.size)\n            if ida_bytes.create_struct(ea, self.size, self.tid):\n                print(\"  Made td at 0x%x: %s\" % (ea, demangled))\n                self.class_name = demangled\n                return\n        print(\"  FAIL :(\")\n        return\n\nclass RTTIClassHierarchyDescriptor(RTTIStruc):\n    bases = None\n\n    msid = get_struc_id(\"RTTIClassHierarchyDescriptor\")\n    if msid != BADADDR:\n        del_struc(msid)\n    msid = add_struc(0xFFFFFFFF, \"RTTIClassHierarchyDescriptor\", False)\n    add_struc_member(msid, \"signature\", BADADDR, FF_DWORD|FF_DATA, -1, 4)\n    add_struc_member(msid, \"attribute\", BADADDR, FF_DWORD|FF_DATA, -1, 4)\n    add_struc_member(msid, \"numBaseClasses\", BADADDR, FF_DWORD|FF_DATA, -1, 4)\n    add_struc_member(msid, \"pBaseClassArray\", BADADDR, FF_DATA|FF_DWORD|FF_0OFF, u.mt_rva().tid, 4)\n    tid = msid\n    struc = get_struc(tid)\n    print(\"Completed Registering RTTIClassHierarchyDescriptor\")\n\n    def __init__(self, ea):\n        print(\"Processing Class Hierarchy Descriptor at 0x%x\" % ea)\n        del_items(ea, DELIT_DELNAMES, get_struc_size(self.tid))\n        if ida_bytes.create_struct(ea, get_struc_size(self.tid), self.tid):\n            baseClasses = get_32bit(ea+get_member_by_name(get_struc(self.tid), \"pBaseClassArray\").soff) + u.x64_imagebase()\n            nb_classes = get_32bit(ea+get_member_by_name(get_struc(self.tid), \"numBaseClasses\").soff)\n            print(\"Baseclasses array at 0x%x\" % baseClasses)\n            # Skip the first base class as it is itself (could check)\n            self.bases = []\n            for i in range(1, nb_classes):\n                baseClass = get_32bit(baseClasses+i*4) + u.x64_imagebase()\n                print(\"base class 0x%x\" % baseClass)\n                ida_bytes.create_dword(baseClasses+i*4, 4)\n                op_offset(baseClasses+i*4, -1, u.REF_OFF|REFINFO_RVA, -1, 0, 0)\n                ida_bytes.create_struct(baseClass, RTTIBaseClassDescriptor.size, RTTIBaseClassDescriptor.tid)\n                typeDescriptor = get_32bit(baseClass) + u.x64_imagebase()\n                self.bases.append(RTTITypeDescriptor(typeDescriptor).class_name)\n\nclass RTTIBaseClassDescriptor(RTTIStruc):\n    msid = get_struc_id(\"RTTIBaseClassDescriptor\")\n    if msid != BADADDR:\n        del_struc(msid)\n    msid = add_struc(0xFFFFFFFF, \"RTTIBaseClassDescriptor\", False)\n    add_struc_member(msid, \"pTypeDescriptor\", BADADDR, FF_DATA|FF_DWORD|FF_0OFF, u.mt_rva().tid, 4)\n    add_struc_member(msid, \"numContainerBases\", BADADDR, FF_DWORD|FF_DATA, -1, 4)\n    add_struc_member(msid, \"PMD\", BADADDR, FF_DATA|FF_DWORD|FF_0OFF, u.mt_rva().tid, 4)\n    add_struc_member(msid, \"attributes\", BADADDR, FF_DWORD|FF_DATA, -1, 4)\n    tid = msid\n    struc = get_struc(tid)\n    size = get_struc_size(tid)\n    print(\"Completed Registering RTTIBaseClassDescriptor\")\n\ndef run_msvc():\n    start = u.rdata.start_ea\n    end = u.rdata.end_ea\n    rdata_size = end-start\n    for offset in range(0, rdata_size-u.PTR_SIZE, u.PTR_SIZE):\n        vtable = start+offset\n        if u.isVtable(vtable):\n            print(\"vtable at : \" + hex(vtable))\n            col = u.get_ptr(vtable-u.PTR_SIZE)\n            if u.within(col, u.valid_ranges):\n                rcol = RTTICompleteObjectLocator(col, vtable)\n    u.add_missing_classes(classes)\n    return classes"
  },
  {
    "path": "readme.md",
    "content": "# SusanRTTI #\n#### Another RTTI Parsing IDA plugin ####\n\n### Features ###\n* All ida-python\n* Class based design, error logging\n* RTTI parsing algorithm scanning for vtables first (instead of\nbruteforcing the entire rdata/data sections)\n* Graphing of class hierarchy (using transitive reduction for clarity)\n* Export functionality to GraphViz (.dot) format\n* Handles RTTI and C++ name demangling for:\n  * X86 GCC\n  * X86 MSVC\n  * X64 GCC\n  * X64 MSVC\n\n### Usage ###\nSimple: First load your binary in IDA. Then run script `classinformerpython.py`\nTo export the dot file, either right click on the class diagram and select\n`export` or just hit F2.\n\n### References ###\nGraphView reference: http://www.graphviz.org/\nOnline viewer: http://www.webgraphviz.com/\n"
  },
  {
    "path": "utils.py",
    "content": "import struct\nimport idaapi\nfrom idc import *\nfrom idaapi import get_segm_by_name, has_xref, get_full_flags, opinfo_t, refinfo_t,\\\nget_32bit, get_64bit, get_imagebase, get_byte\nimport idautils\nfrom idautils import DataRefsTo\n\n# Segments\nwithin = lambda x, rl: any([True for r in rl if r[0]<=x<=r[1]])\n\nclass utils(object):\n    text = 0\n    data = 0\n    rdata = 0\n    valid_ranges = []\n    within = lambda self, x, rl: any([True for r in rl if r[0]<=x<=r[1]])\n\n    REF_OFF = 0\n    x64 = 0\n    PTR_TYPE = 0\n    PTR_SIZE = 0\n\n    def __init__(self):\n        self.text = get_segm_by_name(\".text\")\n        self.data = get_segm_by_name(\".data\")\n        self.rdata = get_segm_by_name(\".rdata\")\n        # try to use rdata if there actually is an rdata segment, otherwise just use data\n        if self.rdata is not None:\n            self.valid_ranges = [(self.rdata.start_ea, self.rdata.end_ea), (self.data.start_ea, self.data.end_ea)]\n        else:\n            self.valid_ranges = [(self.data.start_ea, self.data.end_ea)]\n\n        self.x64 = (idaapi.getseg(here()).bitness == 2)\n        if self.x64:\n            self.PTR_TYPE = FF_QWORD\n            self.REF_OFF = REF_OFF64\n            self.PTR_SIZE = 8\n            self.get_ptr = get_64bit\n        else:\n            self.PTR_TYPE = FF_DWORD\n            self.REF_OFF = REF_OFF32\n            self.PTR_SIZE = 4\n            self.get_ptr = get_32bit\n\n# for 32-bit binaries, the RTTI structs contain absolute addresses, but for\n# 64-bit binaries, they're offsets from the image base.\n    def x64_imagebase(self):\n        if self.x64:\n            return get_imagebase()\n        else:\n            return 0\n\n    def mt_rva(self):\n        ri = refinfo_t()\n        ri.flags = self.REF_OFF\n        ri.target = 0\n        mt = opinfo_t()\n        mt.ri = ri\n        return mt\n\n    def mt_address(self):\n        ri = refinfo_t()\n        ri.flags = self.REF_OFF\n        ri.target = 0\n        mt = opinfo_t()\n        mt.ri = ri\n        return mt\n\n    def mt_ascii(self):\n        ri = refinfo_t()\n        ri.flags = STRTYPE_C\n        ri.target = -1\n        mt = opinfo_t()\n        mt.ri = ri\n        return mt\n\n    def get_strlen(self, addr):\n        strlen = 0\n        while get_byte(addr+strlen) != 0x0 and strlen < 50:\n            strlen+=1\n        #assume no names will ever be longer than 50 bytes\n        if strlen == 50:\n            return None\n        return strlen\n\n    def isVtable(self, addr):\n        function = self.get_ptr(addr)\n        # Check if vtable has ref and its first pointer lies within code segment\n        if has_xref(get_full_flags(addr)) and function >= self.text.start_ea and function <= self.text.end_ea:\n            return True\n        return False\n\n# helper for bin search\n    def ptr_to_bytes(self, val):\n      if self.x64:\n        sv = struct.pack(\"<Q\", val)\n      else:\n        sv = struct.pack(\"<I\", val)\n      return \" \".join(\"%02X\" % ord(c) for c in sv)\n\n    def ptrfirst(self, val):\n      return find_binary(0, SEARCH_CASE|SEARCH_DOWN, self.ptr_to_bytes(val))\n\n    def ptrnext(self, val, ref):\n      return find_binary(ref+1, SEARCH_CASE|SEARCH_DOWN, self.ptr_to_bytes(val))\n\n    def xref_or_find(self, addr, allow_many = False):\n      lrefs = list(DataRefsTo(addr))\n      if len(lrefs) == 0:\n        lrefs = list(idautils.refs(addr, self.ptrfirst, self.ptrnext))\n      if len(lrefs) > 1 and not allow_many:\n          print(\"too many xrefs to %08X\" % addr)\n          return []\n      lrefs = [r for r in lrefs if not is_code(get_full_flags(r))]\n      return lrefs\n\n    def find_string(self, s, afrom=0):\n      print(\"searching for %s\" % s)\n      ea = find_binary(afrom, SEARCH_CASE|SEARCH_DOWN, '\"' + s + '\"')\n      if ea != BADADDR:\n        print(\"Found at %08X\" % ea)\n      return ea\n\n    def ForceDword(self, ea):\n      if ea != BADADDR and ea != 0:\n        if not is_dword(get_full_flags(ea)):\n          del_items(ea, 4, DELIT_SIMPLE)\n          create_data(ea, FF_DWORD, 4, BADADDR)\n        if is_off0(get_full_flags(ea)) and get_fixup_target_type(ea) == -1:\n          # remove the offset\n          op_hex(ea, 0)\n\n    def ForceQword(self, ea):\n      if ea != BADADDR and ea != 0:\n        if not is_qword(get_full_flags(ea)):\n          del_items(ea, 8, DELIT_SIMPLE)\n          create_data(ea, FF_QWORD, 8, BADADDR)\n        if is_off0(get_full_flags(ea)) and get_fixup_target_type(ea) == -1:\n          # remove the offset\n          op_hex(ea, 0)\n\n    def ForcePtr(self, ea, delta = 0):\n      if self.x64:\n        self.ForceQword(ea)\n      else:\n        self.ForceDword(ea)\n      if get_fixup_target_type(ea) != -1 and is_off0(get_full_flags(ea)):\n        # don't touch fixups\n        return\n      pv = self.get_ptr(ea)\n      if pv != 0 and pv != BADADDR:\n        # apply offset again\n        if idaapi.is_spec_ea(pv):\n          delta = 0\n        op_offset(ea, 0, [REF_OFF32, REF_OFF64][self.x64], -1, 0, delta)\n\n# p pointer\n# v vtable pointer (delta ptrsize*2)\n# i integer (32-bit)\n# l integer (32 or 64-bit)\n    def format_struct(self, ea, fmt):\n      for f in fmt:\n        if f in ['p', 'v']:\n          if f == 'v':\n            delta = self.PTR_SIZE*2\n          else:\n            delta = 0\n          self.ForcePtr(ea, delta)\n          ea += self.PTR_SIZE\n        elif f == 'i':\n          self.ForceDword(ea)\n          ea += 4\n        elif f == 'l':\n          if self.x64:\n            self.ForceQword(ea)\n            ea += 8\n          else:\n            self.ForceDword(ea)\n            ea += 4\n      return ea\n\n    def force_name(self, ea, name):\n      if is_tail(get_full_flags(ea)):\n        del_items(ea, 1, DELIT_SIMPLE)\n      set_name(ea, name, SN_NOWARN)\n\n    def is_bad_addr(self, ea):\n      return ea == 0 or ea == BADADDR or idaapi.is_spec_ea(ea) or not is_loaded(ea)\n\n    def vtname(self, name):\n      return \"__ZTV\" + name\n\n# sign extend b low bits in x\n# from \"Bit Twiddling Hacks\"\n    def SIGNEXT(self, x, b):\n        m = 1 << (b - 1)\n        x = x & ((1 << b) - 1)\n        return (x ^ m) - m\n\n    def xref_or_find(self, addr, allow_many = False):\n        lrefs = list(DataRefsTo(addr))\n        if len(lrefs) == 0:\n            lrefs = list(idautils.refs(addr, self.ptrfirst, self.ptrnext))\n        if len(lrefs) > 1 and not allow_many:\n            print(\"too many xrefs to %08X\" % addr)\n            return []\n        lrefs = [r for r in lrefs if not is_code(get_full_flags(r))]\n        return lrefs\n\n    def num2key(self, all_classes):\n        return [k for k in all_classes]\n\n    def add_missing_classes(self, classes):\n        missing = []\n        for c, parents in classes.items():\n            for parent in parents:\n                if parent not in classes.keys():\n                    missing.append(parent)\n        for m in missing:\n            classes[m] = []"
  }
]