improvements

pyclassinformer/method_classifier.py

@@ -1,8 +1,11 @@
 import ida_idaapi
 import ida_funcs
 import ida_name
+import ida_bytes
 import idautils
 
+import pyclassinformer
+
 try:
     ModuleNotFoundError
 except NameError:
@@ -26,6 +29,17 @@
     ida_idaapi.require("pyclassinformer.mc_tree")
     ida_idaapi.require("pyclassinformer.dirtree_utils")
 
+# lazy utils accessor to avoid construction when not available in tests
+_u = None
+def get_utils():
+    global _u
+    if _u is None:
+        try:
+            _u = pyclassinformer.pci_utils.utils()
+        except Exception:
+            _u = None
+    return _u
+
 def change_dir_of_ctors_dtors(paths, data, dirtree):
     path_prefix = "/classes/"
 
@@ -186,6 +200,21 @@ def rename_vfuncs(paths, data):
         rename_funcs(vfunc_eas, class_name.split("<")[0] + "::", is_lib=is_lib)
 
 
+def add_vfunc_index_comments(paths, data):
+    u = pyclassinformer.pci_utils.utils()
+    for vftable_ea in paths:
+        path = paths[vftable_ea]
+        if not path:
+            continue
+
+        vfunc_eas = data[vftable_ea].vfeas
+
+        for index, vfea in enumerate(vfunc_eas):
+            entry_ea = vftable_ea + (index * u.PTR_SIZE)
+            comment = format(index)
+            ida_bytes.set_cmt(entry_ea, comment, 0)
+
+
 def get_base_classes(data):
     paths = {}
     for vftable_ea in data:
@@ -208,7 +237,7 @@ def method_classifier(data, config=None, icon=-1):
         config = pyclassinformer.pci_config.pci_confg()
 
     # check config values to execute or not
-    if not config.exana and not config.mvvm and not config.mvcd and not config.rnvm and not config.rncd:
+    if not config.exana and not config.mvvm and not config.mvcd and not config.rnvm and not config.rncd and not config.rnvi:
         return None
 
     # get base classes
@@ -221,6 +250,10 @@ def method_classifier(data, config=None, icon=-1):
     # rename functions that refer to vftables because they are constructors or destructors
     if config.rncd:
        rename_vftable_ref_funcs(paths, data)
+
+    # add virtual function index comments
+    if config.rnvi:
+        add_vfunc_index_comments(paths, data)
 
     tree = None
     if tree_categorize:
@@ -239,3 +272,4 @@ def method_classifier(data, config=None, icon=-1):
         print("Warning; Your IDA does not have ida_dirtree or find_entry in dirtree_t. Skip creating dirs for classes and moving functions into them.")
 
     return tree
+

pyclassinformer/msvc_rtti.py

@@ -511,18 +511,40 @@ def parse(start, end):
 
         # may be this is a bug on IDA.
         # ida fails to apply a structure type to bytes under some conditions, although create_struct returns True.
-        # to avoid that, apply them again.
+        # to avoid that, apply them again. Cache xref targets per tid to avoid repeated calls.
         ida_auto.auto_wait()
-        #print(len([xrea for xrea in u.get_refs_to(RTTICompleteObjectLocator.tid)]), len([result[x].ea for x in result]))
-        if len([xrea for xrea in u.get_refs_to(RTTICompleteObjectLocator.tid)]) != len([result[x].ea for x in result]):
-            [ida_bytes.create_struct(result[x].ea, RTTICompleteObjectLocator.size, RTTICompleteObjectLocator.tid, True) for x in result]
-        #print(len([xrea for xrea in u.get_refs_to(RTTIClassHierarchyDescriptor.tid)]), len(set([result[x].chd.ea for x in result])))
-        if len([xrea for xrea in u.get_refs_to(RTTIClassHierarchyDescriptor.tid)]) != len(set([result[x].chd.ea for x in result])):
-            [ida_bytes.create_struct(result[x].chd.ea, RTTIClassHierarchyDescriptor.size, RTTIClassHierarchyDescriptor.tid, True) for x in result]
-        #print(len([xrea for xrea in u.get_refs_to(RTTITypeDescriptor.tid)]), len(set([result[x].td.ea for x in result])))
-        if len([xrea for xrea in u.get_refs_to(RTTITypeDescriptor.tid)]) != len(set([result[x].td.ea for x in result])):
-            [ida_bytes.create_struct(result[x].td.ea, result[x].td.size, RTTITypeDescriptor.tid, True) for x in result]
-
+
+        # helper: get xref set for a tid (cache)
+        def get_xref_set(tid):
+            return set(u.get_refs_to(tid))
+
+        # build caches of expected addresses
+        found_col_eas = [result[x].ea for x in result]
+        found_chd_eas = set([result[x].chd.ea for x in result])
+        found_td_eas = set([result[x].td.ea for x in result])
+
+        # apply structs only to addresses missing from IDA's xref results
+        col_xrefs = get_xref_set(RTTICompleteObjectLocator.tid)
+        if len(col_xrefs) != len(found_col_eas):
+            for ea in found_col_eas:
+                if ea not in col_xrefs:
+                    ida_bytes.create_struct(ea, RTTICompleteObjectLocator.size, RTTICompleteObjectLocator.tid, True)
+
+        chd_xrefs = get_xref_set(RTTIClassHierarchyDescriptor.tid)
+        if len(chd_xrefs) != len(found_chd_eas):
+            for ea in found_chd_eas:
+                if ea not in chd_xrefs:
+                    ida_bytes.create_struct(ea, RTTIClassHierarchyDescriptor.size, RTTIClassHierarchyDescriptor.tid, True)
+
+        td_xrefs = get_xref_set(RTTITypeDescriptor.tid)
+        if len(td_xrefs) != len(found_td_eas):
+            for ea in found_td_eas:
+                if ea not in td_xrefs:
+                    # find corresponding td size from result mapping
+                    # result[x].td.size may be different; safe fallback to RTTITypeDescriptor.size
+                    td_size = next((result[x].td.size for x in result if result[x].td.ea == ea), RTTITypeDescriptor.size)
+                    ida_bytes.create_struct(ea, td_size, RTTITypeDescriptor.tid, True)
+
         # for refreshing xrefs to get xrefs from COLs to TDs
         ida_auto.auto_wait()
 

pyclassinformer/pci_config.py

@@ -12,16 +12,18 @@ class pci_config(object):
     mvcd = True
     rnvm = True
     rncd = True
+    rnvi = False
     dirtree = True
 
-    def __init__(self, alldata=False, rtti=True, exana=True, mvvm=True, mvcd=True, rnvm=True, rncd=True):
+    def __init__(self, alldata=False, rtti=True, exana=True, mvvm=True, mvcd=True, rnvm=True, rncd=True, rnvi=True):
         self.alldata = alldata
         self.rtti = rtti
         self.exana = exana
         self.mvvm = mvvm
         self.mvcd = mvcd
         self.rnvm = rnvm
         self.rncd = rncd
+        self.rnvi = rnvi
         self.check_dirtree()
 
     def check_dirtree(self):

pyclassinformer/pci_config_form.py

@@ -21,11 +21,12 @@ def __init__(self, dirtree=True):
 <##Create folders for classes and move virtual methods to them in Functions and Names subviews (IDA 7.7 or later):{mvvm}>
 <##Move functions refer vftables to "possible ctors or dtors" folder under each class folder in Functions and Names subviews (IDA 7.7 or later):{mvcd}>
 <##Rename virtual methods:{rnvm}>
+<##Append virtual index comment:{rnvi}>
 <##Rename possible constructors and destructors:{rncd}>{acts}>
 """, {
             'FormChangeCb': F.FormChangeCb(self.OnFormChange),
             'search_area': F.RadGroupControl(("rdata", "alldata")),
-            'acts': F.ChkGroupControl(("rtti", "exana", "mvvm", "mvcd", "rnvm", "rncd")),
+            'acts': F.ChkGroupControl(("rtti", "exana", "mvvm", "mvcd", "rnvm", "rnvi", "rncd")),
         })
 
         self.dirtree = dirtree
@@ -61,6 +62,7 @@ def set_default_settings(self):
         self.mvvm.checked = True
         self.mvcd.checked = True
         self.rnvm.checked = True
+        self.rnvi.checked = False
         self.rncd.checked = True
 
     @staticmethod
@@ -73,7 +75,7 @@ def show():
         # Execute the form
         ok = f.Execute()
         if ok == 1:
-            pcic = pyclassinformer.pci_config.pci_config(alldata=f.alldata.selected, rtti=f.rtti.checked, exana=f.exana.checked, mvvm=f.mvvm.checked, mvcd=f.mvcd.checked, rnvm=f.rnvm.checked, rncd=f.rncd.checked)
+            pcic = pyclassinformer.pci_config.pci_config(alldata=f.alldata.selected, rtti=f.rtti.checked, exana=f.exana.checked, mvvm=f.mvvm.checked, mvcd=f.mvcd.checked, rnvm=f.rnvm.checked, rnvi=f.rnvi.checked, rncd=f.rncd.checked)
         else:
             return None