PLU-192: Add data types support

pyproject.toml

@@ -29,7 +29,8 @@ dev-dependencies = [
   "pyqt5",
   "pyqt6-stubs",
   "pyqt6",
-  "PySide6"
+  "PySide6",
+  "debugpy>=1.8.17",
 ]
 
 # If you *do* want to install a stub package for IDA later, you can use a sources entry:

reai_toolkit/app/app.py

@@ -14,6 +14,7 @@
 from reai_toolkit.app.services.matching.matching_service import MatchingService
 from reai_toolkit.app.services.rename.rename_service import RenameService
 from reai_toolkit.app.services.upload.upload_service import UploadService
+from reai_toolkit.app.services.data_types.data_types_service import ImportDataTypesService
 
 
 class App:
@@ -40,7 +41,10 @@ def __init__(self, ida_version: str = "UNKNOWN", plugin_version: str = "UNKNOWN"
         self.analysis_status_service = AnalysisStatusService(
             netstore_service=self.netstore_service, sdk_config=sdk_config
         )
-        self.analysis_sync_service = AnalysisSyncService(
+        self.data_types_service = ImportDataTypesService(
+            netstore_service=self.netstore_service, sdk_config=sdk_config
+        )
+        self.analysis_sync_service = AnalysisSyncService(data_types_service=self.data_types_service,
             netstore_service=self.netstore_service, sdk_config=sdk_config
         )
         self.existing_analyses_service = ExistingAnalysesService(
@@ -58,3 +62,4 @@ def __init__(self, ida_version: str = "UNKNOWN", plugin_version: str = "UNKNOWN"
         self.matching_service = MatchingService(
             netstore_service=self.netstore_service, sdk_config=sdk_config
         )
+

reai_toolkit/app/components/dialogs/auto_unstrip_dialog.py

@@ -77,32 +77,37 @@ def _populate_table(self, matches: list[MatchedFunctionSuggestion]) -> None:
         hdr.setStretchLastSection(False)
 
         table.setRowCount(len(matches))
+
+        if QT_VER == 6:
+            flags = QtCore.Qt.ItemFlag.ItemIsSelectable | QtCore.Qt.ItemFlag.ItemIsEnabled
+        else:
+            flags = QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled
 
         for row, m in enumerate(matches):
             # 2. address cell
             addr_item = QtWidgets.QTableWidgetItem(hex(m.function_vaddr))
-            addr_item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
+            addr_item.setFlags(flags)
             addr_item.setToolTip(addr_item.text())  # show on hover
             table.setItem(row, 0, addr_item)
 
             # 3. current name cell
             current_name = get_safe_name(m.function_vaddr) or "<unnamed>"
             cur_item = QtWidgets.QTableWidgetItem(current_name)
-            cur_item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
+            cur_item.setFlags(flags)
             cur_item.setToolTip(current_name)
             table.setItem(row, 1, cur_item)
 
             # 4. suggested name cell
             sug_name = m.suggested_name or ""
             sug_item = QtWidgets.QTableWidgetItem(sug_name)
-            sug_item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
+            sug_item.setFlags(flags)
             sug_item.setToolTip(sug_name)  # tooltip always has full text
             table.setItem(row, 2, sug_item)
 
             # 5. Suggested demangled name cell
             demangled_name = m.suggested_demangled_name or sug_name
             dem_item = QtWidgets.QTableWidgetItem(demangled_name)
-            dem_item.setFlags(QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEnabled)
+            dem_item.setFlags(flags)
             dem_item.setToolTip(demangled_name)  # tooltip always has full text
             table.setItem(row, 3, dem_item)
 

reai_toolkit/app/components/dialogs/matching_dialog.py

@@ -26,6 +26,7 @@
     SummaryEvent,
     ValidFunction,
 )
+from reai_toolkit.app.services.data_types.data_types_service import ImportDataTypesService
 from reai_toolkit.app.services.rename.rename_service import RenameService
 from reai_toolkit.app.services.rename.schema import RenameInput
 
@@ -59,9 +60,10 @@ class MatchingWorker(QtCore.QObject):
     finished = Signal()  # always emitted on exit
     errored = Signal(str)
 
-    def __init__(self, match_service, gen_kwargs: dict):
+    def __init__(self, match_service, data_types_service, gen_kwargs: dict):
         super().__init__()
         self._match_service = match_service
+        self._data_types_service = data_types_service
         self._gen_kwargs = gen_kwargs
         self._stop = False
 
@@ -100,6 +102,7 @@ def __init__(
         func_map: dict[str, int],
         matching_service: MatchingService,
         rename_service: RenameService,
+        data_types_service: ImportDataTypesService,
         parent: QtWidgets.QWidget | None = None,
     ):
         super().__init__(parent=parent)
@@ -108,8 +111,12 @@ def __init__(
 
         self.matching_service = matching_service
         self.rename_service = rename_service
+        self.data_types_service = data_types_service
         self._func_map = func_map
 
+        # Matched function id to original effective address
+        self.matched_func_to_original_ea: dict[int, int] = {}
+
         self.ui = Ui_MatchingPanel()
         self.setWindowTitle("RevEng.AI — Function matching")
         self.ui.setupUi(self)
@@ -214,6 +221,7 @@ def __init__(
 
         if hasattr(self.ui, "okRenameButton"):
             self.ui.okRenameButton.clicked.connect(self.enqueue_renames)
+            self.ui.okRenameButton.clicked.connect(self.import_data_types)
 
         # ----------------- Util buttons -----------------
         self.ui.btnClearSelection.clicked.connect(
@@ -699,7 +707,7 @@ def start_ann(self):
         self.stop_ann()  # ensure previous worker is cleaned up
 
         self._matching_thread = QtCore.QThread(self)
-        self._matching_worker = MatchingWorker(self.matching_service, gen_kwargs)
+        self._matching_worker = MatchingWorker(self.matching_service, self.data_types_service, gen_kwargs)
         self._matching_worker.moveToThread(self._matching_thread)
 
         # connections
@@ -1083,6 +1091,8 @@ def display_matching_results_multiple_functions(self, query: str = ""):
                 r.matched_functions[0] if r.matched_functions else None
             )
 
+            self.matched_func_to_original_ea[matched_function.function_id] = self._func_map[str(r.function_id)]
+
             # Column 3: Matched Name
             table.setItem(
                 row,
@@ -1183,6 +1193,9 @@ def enqueue_renames(self):
 
         except Exception as e:
             print(f"Failed to enqueue renames: {e}")
+
+    def import_data_types(self):
+        self.data_types_service.import_data_types(self.matched_func_to_original_ea)
 
     # =====================================================================
     # (Optional) page-switch helpers

reai_toolkit/app/coordinator.py

@@ -73,6 +73,7 @@ def __init__(self, app, factory, log):
             log=log,
             auto_unstrip_service=app.auto_unstrip_service,
             rename_service=app.rename_service,
+            data_types_service=app.data_types_service
         )
 
         self.ai_decompc: AiDecompCoordinator = AiDecompCoordinator(

reai_toolkit/app/coordinators/auto_unstrip_coordinator.py

@@ -10,6 +10,7 @@
 )
 from reai_toolkit.app.services.rename.rename_service import RenameService
 from reai_toolkit.app.services.rename.schema import RenameInput
+from reai_toolkit.app.services.data_types.data_types_service import ImportDataTypesService
 
 
 class AutoUnstripCoordinator(BaseCoordinator):
@@ -25,10 +26,12 @@ def __init__(
         log,
         auto_unstrip_service: AutoUnstripService,
         rename_service: RenameService,
+        data_types_service: ImportDataTypesService,
     ):
         super().__init__(app=app, factory=factory, log=log)
         self.auto_unstrip_service = auto_unstrip_service
         self.rename_service = rename_service
+        self.data_types_service: ImportDataTypesService = data_types_service
 
     def run_dialog(self) -> None:
         if self.auto_unstrip_service.is_worker_running():
@@ -46,14 +49,15 @@ def run_dialog(self) -> None:
     def _open_auto_unstrip_dialog(self) -> None:
         self.factory.auto_unstrip(response=self.last_response.data).open_modal()
 
-    def _on_complete(self, response: GenericApiReturn[AutoUnstripResponse]):
+    def _on_complete(self, response: GenericApiReturn[AutoUnstripResponse]) -> None:
         print("Auto-unstrip process completed.")
 
         if not response.success:
             self.safe_error(message=response.error_message)
             return
 
         rename_list = []
+        data_types_mapping: dict[int, int] = {}
 
         self.last_response = response
 
@@ -65,8 +69,10 @@ def _on_complete(self, response: GenericApiReturn[AutoUnstripResponse]):
                     new_name=function.suggested_demangled_name,
                 )
             )
+            data_types_mapping[function.function_id] = function.function_vaddr
 
         self.rename_service.enqueue_rename(rename_list=rename_list)
+        self.data_types_service.import_data_types(data_types_mapping)
 
         ida_kernwin.execute_ui_requests([self._open_auto_unstrip_dialog])
 

reai_toolkit/app/coordinators/sync_analysis_coordinator.py

@@ -48,9 +48,8 @@ def _on_complete(
         """
         if generic_return.success:
             self.safe_info(
-                msg=f"Analysis data synced successfully. \n\nSynced {generic_return.data.matched_local_function_count} functions with remote analysis."
-                + f"\n{generic_return.data.unmatched_local_function_count} local functions not present in remote analysis."
-                + f"\n{generic_return.data.unmatched_remote_function_count} remote functions not present in local analysis."
+                msg=f"Analysis data synced successfully. \n\nSynced {generic_return.data.matched_function_count} functions with remote analysis."
+                + f"\n{generic_return.data.unmatched_function_count} local functions not present in remote analysis."
             )
         else:
             self.safe_error(message=generic_return.error_message)

reai_toolkit/app/factory.py

@@ -70,6 +70,7 @@ def function_matching(
             func_map=func_map,
             matching_service=self.app.matching_service,
             rename_service=self.app.rename_service,
+            data_types_service=self.app.data_types_service,
             parent=self.parent,
         )