diff --git a/reai_toolkit/app/components/dialogs/matching_dialog.py b/reai_toolkit/app/components/dialogs/matching_dialog.py index 9c40754b..f739205d 100644 --- a/reai_toolkit/app/components/dialogs/matching_dialog.py +++ b/reai_toolkit/app/components/dialogs/matching_dialog.py @@ -41,7 +41,7 @@ Ui_MatchingPanel, ) -print("[AnnDialog] Qt version:", QT_VER) +logger.debug(f"[AnnDialog] Qt version: {QT_VER}") DEBOUNCE_MS = 250 @@ -218,7 +218,7 @@ def __init__( self._progress_bar = self.ui.loadingBar if hasattr(self.ui, "loadingBar") else None if self._progress_bar is not None: - print("[AnnDialog] Progress bar found") + logger.debug("[AnnDialog] Progress bar found") self._progress_bar.setMinimum(0) self._progress_bar.setMaximum(0) @@ -705,7 +705,7 @@ def start_ann(self): debug_all=True if self.ui.chkDebugSymbols.isChecked() else None, ) - print(f"[AnnDialog] Starting ANN with args: {gen_kwargs}") + logger.debug(f"[AnnDialog] Starting ANN with args: {gen_kwargs}") self.stop_ann() # ensure previous worker is cleaned up @@ -732,7 +732,7 @@ def on_ann_error(self, msg: str): if self._progress_bar: self._progress_bar.setVisible(False) self.toggle_search_ui(pause=False) - print(f"[AnnDialog] ANN worker error: {msg}") + logger.error(f"[AnnDialog] ANN worker error: {msg}") self.error_dialog = ErrorDialog(error_message=msg, parent=self) self.error_dialog.show() @@ -746,7 +746,7 @@ def stop_ann(self): try: self._matching_worker.stop() except Exception as e: - print(f"[AnnDialog] Failed to stop ANN worker: {e}") + logger.debug(f"[AnnDialog] Failed to stop ANN worker: {e}") pass if self._matching_thread is not None: # Let it exit naturally; Qt handles cleanup via the signals we set @@ -763,7 +763,7 @@ def on_ann_event(self, ev): etype = getattr(ev, "event", None) etype: MatchEvent - print( + logger.debug( f"[AnnDialog] ANN event: {etype}, completed: {getattr(ev, 'completed', None)}/{getattr(ev, 'total', None)}, ok: {getattr(ev, 'ok', None)}" ) @@ -978,7 +978,7 @@ def display_matching_results_single_function(self, query: str = ""): table.setWordWrap(True) table.resizeRowsToContents() - print( + logger.debug( f"Single: Displayed {len(self.matching_results.results[0].matched_functions)} ANN results" ) @@ -1202,7 +1202,7 @@ def select_all() -> None: self.ui.btnResultClearAll.clicked.connect(clear_all) self.ui.btnResultSelectAll.clicked.connect(select_all) - print(f"Multiple: Displayed {len(self.matching_results.results)} ANN results") + logger.debug(f"Multiple: Displayed {len(self.matching_results.results)} ANN results") def enqueue_renames(self): # Read checked rows from results table @@ -1216,7 +1216,7 @@ def enqueue_renames(self): function_id = item.data(QtCore.Qt.ItemDataRole.UserRole) matched_item = table.item(r, MatchColumns.MATCHED_NAME).text() vaddr = self.rename_service.function_id_to_vaddr(function_id) - print(f"RENAME: Function ID {function_id} -> {matched_item}") + logger.debug(f"RENAME: Function ID {function_id} -> {matched_item}") rename_list.append( RenameInput( function_id=function_id, @@ -1228,7 +1228,7 @@ def enqueue_renames(self): self.accept() except Exception as e: - print(f"Failed to enqueue renames: {e}") + logger.error(f"Failed to enqueue renames: {e}") def import_data_types(self): selected_matches: dict[int, int] = {} diff --git a/reai_toolkit/app/components/forms/ai_decomp/__init__.py b/reai_toolkit/app/components/forms/ai_decomp/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_qt5.ui b/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_qt5.ui deleted file mode 100644 index 1e31bf42..00000000 --- a/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_qt5.ui +++ /dev/null @@ -1,279 +0,0 @@ - - - AiDecompPanel - - - - 0 - 0 - 900 - 600 - - - - - 900 - 600 - - - - - 900 - 600 - - - - RevEng.AI — Ai Decomp - - - - QGroupBox{font-weight:600;margin-top:8px;} - QGroupBox::title{subcontrol-origin: margin; left: 6px; padding:2px 4px;} - QLineEdit{border:1px solid #3c3f44;border-radius:6px;padding:6px;} - QLineEdit:focus{border-color:#5865f2;} - QPushButton{border:1px solid #3c3f44;border-radius:8px;padding:6px 12px;} - QPushButton:hover{border-color:#5865f2;} - #okButton{background:rgb(35,144,236);color:white;border:0;border-radius:8px;padding:6px 12px;} - #cancelButton{background:rgb(153,12,25);color:white;border:0;border-radius:8px;padding:6px 12px;} - - /* thumbs styling (SVGs must use fill="currentColor") */ - QPushButton#buttonThumbsUp { color: palette(text); } - QPushButton#buttonThumbsUp:checked { color: #16a34a; } /* green */ - - QPushButton#buttonThumbsDown { color: palette(text); } - QPushButton#buttonThumbsDown:checked { color: #dc2626; } /* red */ - - - - - 10 - - - 12 - - - 12 - - - 12 - - - 12 - - - - - QFrame::NoFrame - - - - 12 - - - - - - 72 - 72 - - - - - 96 - 96 - - - - QFrame::NoFrame - - - - - - ../../resources/reveng_ai_logo.jpg - - - true - - - - - - - - - font-size:18px;font-weight:600; - - - RevEng.AI — AI Decomp - - - - - - - false - - - color: palette(text); - - - AI Decompilation - - - - - - - - - Qt::Horizontal - - - - 40 - 20 - - - - - - - - - - - - color: palette(text); - - - Rate the AI-Decompilation - - - - - - - - - - 60 - 0 - - - - Thumbs down - - - - - - - 24 - 24 - - - - true - - - true - - - - - - - - 60 - 0 - - - - Thumbs up - - - - - - - ../../resources/thumb_up.svg../../resources/thumb_up.svg - - - - 24 - 24 - - - - true - - - true - - - - - - - Qt::Horizontal - - - - 20 - 20 - - - - - - - - - - - - - - - - - - 0 - 150 - - - - Log - - - - 6 - - - - - - 0 - 0 - - - - QPlainTextEdit::NoWrap - - - true - - - - - - - - - - - - - - diff --git a/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_ui_uic5.py b/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_ui_uic5.py deleted file mode 100644 index 6ef95045..00000000 --- a/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_ui_uic5.py +++ /dev/null @@ -1,162 +0,0 @@ -# -*- coding: utf-8 -*- - -# Form implementation generated from reading ui file 'reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_qt5.ui' -# -# Created by: PyQt5 UI code generator 5.15.11 -# -# WARNING: Any manual changes made to this file will be lost when pyuic5 is -# run again. Do not edit this file unless you know what you are doing. - - -from PyQt5 import QtCore, QtGui, QtWidgets - - -class Ui_AiDecompPanel(object): - def setupUi(self, AiDecompPanel): - AiDecompPanel.setObjectName("AiDecompPanel") - AiDecompPanel.resize(900, 800) - AiDecompPanel.setMinimumSize(QtCore.QSize(900, 800)) - AiDecompPanel.setMaximumSize(QtCore.QSize(900, 800)) - AiDecompPanel.setStyleSheet( - "\n" - " QGroupBox{font-weight:600;margin-top:8px;}\n" - " QGroupBox::title{subcontrol-origin: margin; left: 6px; padding:2px 4px;}\n" - " QLineEdit{border:1px solid #3c3f44;border-radius:6px;padding:6px;}\n" - " QLineEdit:focus{border-color:#5865f2;}\n" - " QPushButton{border:1px solid #3c3f44;border-radius:8px;padding:6px 12px;}\n" - " QPushButton:hover{border-color:#5865f2;}\n" - " #okButton{background:rgb(35,144,236);color:white;border:0;border-radius:8px;padding:6px 12px;}\n" - " #cancelButton{background:rgb(153,12,25);color:white;border:0;border-radius:8px;padding:6px 12px;}\n" - "\n" - " /* thumbs styling (requires SVGs that use currentColor) */\n" - " QPushButton#buttonThumbsUp { color: palette(text); }\n" - " QPushButton#buttonThumbsUp:checked { color: #16a34a; } /* green */\n" - "\n" - " QPushButton#buttonThumbsDown { color: palette(text); }\n" - " QPushButton#buttonThumbsDown:checked { color: #dc2626; } /* red */\n" - " " - ) - self.rootLayout = QtWidgets.QVBoxLayout(AiDecompPanel) - self.rootLayout.setContentsMargins(12, 12, 12, 12) - self.rootLayout.setSpacing(10) - self.rootLayout.setObjectName("rootLayout") - self.header = QtWidgets.QFrame(AiDecompPanel) - self.header.setFrameShape(QtWidgets.QFrame.NoFrame) - self.header.setObjectName("header") - self.headerLayout = QtWidgets.QHBoxLayout(self.header) - self.headerLayout.setSpacing(12) - self.headerLayout.setObjectName("headerLayout") - self.logoArea = QtWidgets.QLabel(self.header) - self.logoArea.setMinimumSize(QtCore.QSize(72, 72)) - self.logoArea.setMaximumSize(QtCore.QSize(96, 96)) - self.logoArea.setFrameShape(QtWidgets.QFrame.NoFrame) - self.logoArea.setText("") - self.logoArea.setPixmap( - QtGui.QPixmap( - "reai_toolkit/app/components/forms/ai_decomp/../../resources/reveng_ai_logo.jpg" - ) - ) - self.logoArea.setScaledContents(True) - self.logoArea.setObjectName("logoArea") - self.headerLayout.addWidget(self.logoArea) - self.titleLayout = QtWidgets.QVBoxLayout() - self.titleLayout.setObjectName("titleLayout") - self.titleLabel = QtWidgets.QLabel(self.header) - self.titleLabel.setStyleSheet("font-size:18px;font-weight:600;") - self.titleLabel.setObjectName("titleLabel") - self.titleLayout.addWidget(self.titleLabel) - self.subtitleLabel = QtWidgets.QLabel(self.header) - self.subtitleLabel.setAutoFillBackground(False) - self.subtitleLabel.setStyleSheet("color: palette(text);") - self.subtitleLabel.setObjectName("subtitleLabel") - self.titleLayout.addWidget(self.subtitleLabel) - self.headerLayout.addLayout(self.titleLayout) - spacerItem = QtWidgets.QSpacerItem( - 40, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum - ) - self.headerLayout.addItem(spacerItem) - self.ratingLayout = QtWidgets.QHBoxLayout() - self.ratingLayout.setObjectName("ratingLayout") - self.verticalLayout = QtWidgets.QVBoxLayout() - self.verticalLayout.setObjectName("verticalLayout") - self.label = QtWidgets.QLabel(self.header) - self.label.setStyleSheet("color: palette(text);") - self.label.setObjectName("label") - self.verticalLayout.addWidget(self.label) - self.horizontalLayout_2 = QtWidgets.QHBoxLayout() - self.horizontalLayout_2.setObjectName("horizontalLayout_2") - self.buttonThumbsDown = QtWidgets.QPushButton(self.header) - self.buttonThumbsDown.setMinimumSize(QtCore.QSize(60, 0)) - self.buttonThumbsDown.setText("") - icon = QtGui.QIcon() - icon.addPixmap( - QtGui.QPixmap( - "reai_toolkit/app/components/forms/ai_decomp/../../resources/thumb_down.svg" - ), - QtGui.QIcon.Normal, - QtGui.QIcon.Off, - ) - self.buttonThumbsDown.setIcon(icon) - self.buttonThumbsDown.setIconSize(QtCore.QSize(24, 24)) - self.buttonThumbsDown.setCheckable(True) - self.buttonThumbsDown.setObjectName("buttonThumbsDown") - self.horizontalLayout_2.addWidget(self.buttonThumbsDown) - self.buttonThumbsUp = QtWidgets.QPushButton(self.header) - self.buttonThumbsUp.setMinimumSize(QtCore.QSize(60, 0)) - self.buttonThumbsUp.setText("") - icon1 = QtGui.QIcon() - icon1.addPixmap( - QtGui.QPixmap( - "reai_toolkit/app/components/forms/ai_decomp/../../resources/thumb_up.svg" - ), - QtGui.QIcon.Normal, - QtGui.QIcon.Off, - ) - self.buttonThumbsUp.setIcon(icon1) - self.buttonThumbsUp.setIconSize(QtCore.QSize(24, 24)) - self.buttonThumbsUp.setCheckable(True) - self.buttonThumbsUp.setObjectName("buttonThumbsUp") - self.horizontalLayout_2.addWidget(self.buttonThumbsUp) - spacerItem1 = QtWidgets.QSpacerItem( - 20, 20, QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Minimum - ) - self.horizontalLayout_2.addItem(spacerItem1) - self.verticalLayout.addLayout(self.horizontalLayout_2) - self.ratingLayout.addLayout(self.verticalLayout) - self.headerLayout.addLayout(self.ratingLayout) - self.rootLayout.addWidget(self.header) - self.groupEndpoints = QtWidgets.QGroupBox(AiDecompPanel) - self.groupEndpoints.setMinimumSize(QtCore.QSize(0, 150)) - self.groupEndpoints.setObjectName("groupEndpoints") - self.errorLayout = QtWidgets.QVBoxLayout(self.groupEndpoints) - self.errorLayout.setSpacing(6) - self.errorLayout.setObjectName("errorLayout") - self.codeEditor = QtWidgets.QPlainTextEdit(self.groupEndpoints) - sizePolicy = QtWidgets.QSizePolicy( - QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding - ) - sizePolicy.setHorizontalStretch(0) - sizePolicy.setVerticalStretch(0) - sizePolicy.setHeightForWidth(self.codeEditor.sizePolicy().hasHeightForWidth()) - self.codeEditor.setSizePolicy(sizePolicy) - self.codeEditor.setLineWrapMode(QtWidgets.QPlainTextEdit.NoWrap) - self.codeEditor.setReadOnly(True) - self.codeEditor.setObjectName("codeEditor") - self.errorLayout.addWidget(self.codeEditor) - self.rootLayout.addWidget(self.groupEndpoints) - self.linkRow = QtWidgets.QHBoxLayout() - self.linkRow.setObjectName("linkRow") - self.rootLayout.addLayout(self.linkRow) - - self.retranslateUi(AiDecompPanel) - QtCore.QMetaObject.connectSlotsByName(AiDecompPanel) - - def retranslateUi(self, AiDecompPanel): - _translate = QtCore.QCoreApplication.translate - AiDecompPanel.setWindowTitle( - _translate("AiDecompPanel", "RevEng.AI — Ai Decomp") - ) - self.titleLabel.setText(_translate("AiDecompPanel", "RevEng.AI — AI Decomp")) - self.subtitleLabel.setText(_translate("AiDecompPanel", "AI Decompilation")) - self.label.setText(_translate("AiDecompPanel", "Rate the AI-Decompilation")) - self.groupEndpoints.setTitle(_translate("AiDecompPanel", "Log")) diff --git a/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_ui_uic6.py b/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_ui_uic6.py deleted file mode 100644 index b34f440e..00000000 --- a/reai_toolkit/app/components/forms/ai_decomp/ai_decomp_panel_ui_uic6.py +++ /dev/null @@ -1,216 +0,0 @@ -# -*- coding: utf-8 -*- - -################################################################################ -## Form generated from reading UI file 'ai_decomp_panel_qt5.ui' -## -## Created by: Qt User Interface Compiler version 6.9.1 -## -## WARNING! All changes made in this file will be lost when recompiling UI file! -################################################################################ - -from PySide6.QtCore import ( - QCoreApplication, - QMetaObject, - QSize, -) -from PySide6.QtGui import ( - QIcon, - QPixmap, -) -from PySide6.QtWidgets import ( - QFrame, - QGroupBox, - QHBoxLayout, - QLabel, - QPlainTextEdit, - QPushButton, - QSizePolicy, - QSpacerItem, - QVBoxLayout, -) - - -class Ui_AiDecompPanel(object): - def setupUi(self, AiDecompPanel): - if not AiDecompPanel.objectName(): - AiDecompPanel.setObjectName("AiDecompPanel") - AiDecompPanel.resize(900, 800) - AiDecompPanel.setMinimumSize(QSize(900, 800)) - AiDecompPanel.setMaximumSize(QSize(900, 800)) - AiDecompPanel.setStyleSheet( - "\n" - " QGroupBox{font-weight:600;margin-top:8px;}\n" - " QGroupBox::title{subcontrol-origin: margin; left: 6px; padding:2px 4px;}\n" - " QLineEdit{border:1px solid #3c3f44;border-radius:6px;padding:6px;}\n" - " QLineEdit:focus{border-color:#5865f2;}\n" - " QPushButton{border:1px solid #3c3f44;border-radius:8px;padding:6px 12px;}\n" - " QPushButton:hover{border-color:#5865f2;}\n" - " #okButton{background:rgb(35,144,236);color:white;border:0;border-radius:8px;padding:6px 12px;}\n" - " #cancelButton{background:rgb(153,12,25);color:white;border:0;border-radius:8px;padding:6px 12px;}\n" - "\n" - " /* thumbs styling (requires SVGs that use currentColor) */\n" - " QPushButton#buttonThumbsUp { color: palette(text); }\n" - " QPushButton#buttonThumbsUp:checked { color: #16a34a; } /* green */\n" - "\n" - " QPushButton#buttonThumbsDown { color: palette(text); }\n" - " QPushButton#buttonThumbsDown:checked { color: #dc2626; } /* red */\n" - " " - ) - self.rootLayout = QVBoxLayout(AiDecompPanel) - self.rootLayout.setSpacing(10) - self.rootLayout.setObjectName("rootLayout") - self.rootLayout.setContentsMargins(12, 12, 12, 12) - self.header = QFrame(AiDecompPanel) - self.header.setObjectName("header") - self.header.setFrameShape(QFrame.NoFrame) - self.headerLayout = QHBoxLayout(self.header) - self.headerLayout.setSpacing(12) - self.headerLayout.setObjectName("headerLayout") - self.logoArea = QLabel(self.header) - self.logoArea.setObjectName("logoArea") - self.logoArea.setMinimumSize(QSize(72, 72)) - self.logoArea.setMaximumSize(QSize(96, 96)) - self.logoArea.setFrameShape(QFrame.NoFrame) - self.logoArea.setPixmap(QPixmap("../../resources/reveng_ai_logo.jpg")) - self.logoArea.setScaledContents(True) - - self.headerLayout.addWidget(self.logoArea) - - self.titleLayout = QVBoxLayout() - self.titleLayout.setObjectName("titleLayout") - self.titleLabel = QLabel(self.header) - self.titleLabel.setObjectName("titleLabel") - self.titleLabel.setStyleSheet("font-size:18px;font-weight:600;") - - self.titleLayout.addWidget(self.titleLabel) - - self.subtitleLabel = QLabel(self.header) - self.subtitleLabel.setObjectName("subtitleLabel") - self.subtitleLabel.setAutoFillBackground(False) - self.subtitleLabel.setStyleSheet("color: palette(text);") - - self.titleLayout.addWidget(self.subtitleLabel) - - self.headerLayout.addLayout(self.titleLayout) - - self.horizontalSpacer_2 = QSpacerItem( - 40, 20, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Minimum - ) - - self.headerLayout.addItem(self.horizontalSpacer_2) - - self.ratingLayout = QHBoxLayout() - self.ratingLayout.setObjectName("ratingLayout") - self.verticalLayout = QVBoxLayout() - self.verticalLayout.setObjectName("verticalLayout") - self.label = QLabel(self.header) - self.label.setObjectName("label") - self.label.setStyleSheet("color: palette(text);") - - self.verticalLayout.addWidget(self.label) - - self.horizontalLayout_2 = QHBoxLayout() - self.horizontalLayout_2.setObjectName("horizontalLayout_2") - self.buttonThumbsDown = QPushButton(self.header) - self.buttonThumbsDown.setObjectName("buttonThumbsDown") - self.buttonThumbsDown.setMinimumSize(QSize(60, 0)) - icon = QIcon() - icon.addFile( - "../../resources/thumb_down.svg", - QSize(), - QIcon.Mode.Normal, - QIcon.State.Off, - ) - self.buttonThumbsDown.setIcon(icon) - self.buttonThumbsDown.setIconSize(QSize(24, 24)) - self.buttonThumbsDown.setCheckable(True) - - self.horizontalLayout_2.addWidget(self.buttonThumbsDown) - - self.buttonThumbsUp = QPushButton(self.header) - self.buttonThumbsUp.setObjectName("buttonThumbsUp") - self.buttonThumbsUp.setMinimumSize(QSize(60, 0)) - icon1 = QIcon() - icon1.addFile( - "../../resources/thumb_up.svg", QSize(), QIcon.Mode.Normal, QIcon.State.Off - ) - self.buttonThumbsUp.setIcon(icon1) - self.buttonThumbsUp.setIconSize(QSize(24, 24)) - self.buttonThumbsUp.setCheckable(True) - - self.horizontalLayout_2.addWidget(self.buttonThumbsUp) - - self.horizontalSpacer = QSpacerItem( - 20, 20, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Minimum - ) - - self.horizontalLayout_2.addItem(self.horizontalSpacer) - - self.verticalLayout.addLayout(self.horizontalLayout_2) - - self.ratingLayout.addLayout(self.verticalLayout) - - self.headerLayout.addLayout(self.ratingLayout) - - self.rootLayout.addWidget(self.header) - - self.groupEndpoints = QGroupBox(AiDecompPanel) - self.groupEndpoints.setObjectName("groupEndpoints") - self.groupEndpoints.setMinimumSize(QSize(0, 150)) - self.errorLayout = QVBoxLayout(self.groupEndpoints) - self.errorLayout.setSpacing(6) - self.errorLayout.setObjectName("errorLayout") - self.codeEditor = QPlainTextEdit(self.groupEndpoints) - self.codeEditor.setObjectName("codeEditor") - sizePolicy = QSizePolicy( - QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding - ) - sizePolicy.setHorizontalStretch(0) - sizePolicy.setVerticalStretch(0) - sizePolicy.setHeightForWidth(self.codeEditor.sizePolicy().hasHeightForWidth()) - self.codeEditor.setSizePolicy(sizePolicy) - self.codeEditor.setLineWrapMode(QPlainTextEdit.NoWrap) - self.codeEditor.setReadOnly(True) - - self.errorLayout.addWidget(self.codeEditor) - - self.rootLayout.addWidget(self.groupEndpoints) - - self.linkRow = QHBoxLayout() - self.linkRow.setObjectName("linkRow") - - self.rootLayout.addLayout(self.linkRow) - - self.retranslateUi(AiDecompPanel) - - QMetaObject.connectSlotsByName(AiDecompPanel) - - # setupUi - - def retranslateUi(self, AiDecompPanel): - AiDecompPanel.setWindowTitle( - QCoreApplication.translate( - "AiDecompPanel", "RevEng.AI \u2014 Ai Decomp", None - ) - ) - self.logoArea.setText("") - self.titleLabel.setText( - QCoreApplication.translate( - "AiDecompPanel", "RevEng.AI \u2014 AI Decomp", None - ) - ) - self.subtitleLabel.setText( - QCoreApplication.translate("AiDecompPanel", "AI Decompilation", None) - ) - self.label.setText( - QCoreApplication.translate( - "AiDecompPanel", "Rate the AI-Decompilation", None - ) - ) - self.buttonThumbsDown.setText("") - self.buttonThumbsUp.setText("") - self.groupEndpoints.setTitle( - QCoreApplication.translate("AiDecompPanel", "Log", None) - ) - - # retranslateUi diff --git a/reai_toolkit/app/components/tabs/ai_decomp_tab.py b/reai_toolkit/app/components/tabs/ai_decomp_tab.py index 256d8ef4..54db1fe7 100644 --- a/reai_toolkit/app/components/tabs/ai_decomp_tab.py +++ b/reai_toolkit/app/components/tabs/ai_decomp_tab.py @@ -4,7 +4,52 @@ import ida_kernwin as kw from libbs.decompilers.ida.compat import execute_ui -from reai_toolkit.app.core.qt_compat import QtCore, QtGui, QtWidgets +from reai_toolkit.app.core.qt_compat import QtCore, QtGui, QtWidgets, Signal + + +_WORD_UNDER_CURSOR = getattr( + getattr(QtGui.QTextCursor, "SelectionType", QtGui.QTextCursor), "WordUnderCursor" +) + + +def _menu_exec(menu, pos): + fn = getattr(menu, "exec_", None) or getattr(menu, "exec") + return fn(pos) + + +class _DecompEditor(QtWidgets.QPlainTextEdit): + renameRequested = Signal(int, str) + commentEditRequested = Signal(int) + commentRemoveRequested = Signal(int) + + def mouseDoubleClickEvent(self, event) -> None: + super().mouseDoubleClickEvent(event) + cursor = self.textCursor() + word = cursor.selectedText() + if word: + self.renameRequested.emit(cursor.blockNumber(), word) + + def contextMenuEvent(self, event) -> None: + cursor = self.cursorForPosition(event.pos()) + line = cursor.blockNumber() + cursor.select(_WORD_UNDER_CURSOR) + word = cursor.selectedText() + + menu = self.createStandardContextMenu() + menu.addSeparator() + act_rename = menu.addAction(f"Rename '{word}'…") if word else None + act_comment = menu.addAction("Add / edit comment…") + act_remove = menu.addAction("Remove comment") + + chosen = _menu_exec(menu, event.globalPos()) + if chosen is None: + return + if act_rename is not None and chosen == act_rename: + self.renameRequested.emit(line, word) + elif chosen == act_comment: + self.commentEditRequested.emit(line) + elif chosen == act_remove: + self.commentRemoveRequested.emit(line) class AIDecompView(kw.PluginForm): @@ -22,8 +67,13 @@ class AIDecompView(kw.PluginForm): def __init__(self, on_closed: Optional[Callable[[], None]] = None) -> None: super().__init__() self._on_closed: Callable[[], None] | None = on_closed + self.on_refresh: Callable[[], None] | None = None + self.on_rename: Callable[[int, str], None] | None = None + self.on_edit_comment: Callable[[int], None] | None = None + self.on_remove_comment: Callable[[int], None] | None = None self._parent_window: QtWidgets.QWidget | None = None - self._editor: QtWidgets.QPlainTextEdit | None = None + self._editor: _DecompEditor | None = None + self._refresh_btn: QtWidgets.QPushButton | None = None self._highlighter: CppHighlighter | None = None def Create(self, title: Any) -> Any: @@ -52,10 +102,23 @@ def OnCreate(self, form) -> None: layout = QtWidgets.QVBoxLayout(self._parent_window) layout.setContentsMargins(0, 0, 0, 0) + # Header + header = QtWidgets.QHBoxLayout() + title = QtWidgets.QLabel("RevEng.AI — AI Decomp", self._parent_window) + header.addWidget(title) + header.addStretch(1) + self._refresh_btn = QtWidgets.QPushButton("Refresh", self._parent_window) + self._refresh_btn.clicked.connect(self._on_refresh_clicked) + header.addWidget(self._refresh_btn) + layout.addLayout(header) + # Editor - self._editor = QtWidgets.QPlainTextEdit(self._parent_window) + self._editor = _DecompEditor(self._parent_window) self._editor.setReadOnly(True) self._editor.setLineWrapMode(QtWidgets.QPlainTextEdit.NoWrap) + self._editor.renameRequested.connect(self._on_rename_requested) + self._editor.commentEditRequested.connect(self._on_edit_comment_requested) + self._editor.commentRemoveRequested.connect(self._on_remove_comment_requested) # Monospace font tuned for IDA font = QtGui.QFont( @@ -83,8 +146,25 @@ def OnClose(self, form) -> None: logger.warning(f"on_closed callback failed: {e}") self._highlighter = None self._editor = None + self._refresh_btn = None self._parent_window = None + def _on_refresh_clicked(self) -> None: + if self.on_refresh: + self.on_refresh() + + def _on_rename_requested(self, line: int, word: str) -> None: + if self.on_rename: + self.on_rename(line, word) + + def _on_edit_comment_requested(self, line: int) -> None: + if self.on_edit_comment: + self.on_edit_comment(line) + + def _on_remove_comment_requested(self, line: int) -> None: + if self.on_remove_comment: + self.on_remove_comment(line) + # --- public API ------------------------------------------------ @execute_ui def update_view_content(self, code: str) -> None: @@ -97,8 +177,6 @@ def update_view_content(self, code: str) -> None: finally: self._editor.blockSignals(False) - - def clear(self) -> None: self.update_view_content("") diff --git a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py index a04fbf8c..f37ce760 100644 --- a/reai_toolkit/app/coordinators/ai_decomp_coordinator.py +++ b/reai_toolkit/app/coordinators/ai_decomp_coordinator.py @@ -5,11 +5,17 @@ from libbs.decompilers.ida.compat import execute_read, execute_ui from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData -from revengai.models.inline_comment import InlineComment from revengai.models.summary_data import SummaryData +from revengai.models.tokenised_data import TokenisedData from reai_toolkit.app.app import App from reai_toolkit.app.components.tabs.ai_decomp_tab import AIDecompView +from reai_toolkit.app.coordinators.ai_decomp_render import ( + RenderModel, + index_of_identifier, + render_view_with_map, + resolve_token, +) from reai_toolkit.app.coordinators.base_coordinator import BaseCoordinator from reai_toolkit.app.core.shared_schema import GenericApiReturn from reai_toolkit.app.factory import DialogFactory @@ -34,6 +40,8 @@ def __init__( self._current_decomp: DecompilationData | None = None self._current_summary: SummaryData | None = None self._current_comments: CommentsData | None = None + self._current_tokenised: TokenisedData | None = None + self._baseline: RenderModel | None = None def enable_function_tracking(self) -> None: if self._decomp_hooks is None: @@ -56,12 +64,18 @@ def ensure_tracking(self) -> None: def run_dialog(self) -> None: if self._decomp_view is None: self._decomp_view = self.factory.ai_decomp(on_closed=self._on_pane_closed) + self._decomp_view.on_refresh = self.refresh_current + self._decomp_view.on_rename = self.request_rename + self._decomp_view.on_edit_comment = self.request_edit_comment + self._decomp_view.on_remove_comment = self.request_remove_comment self._decomp_view.Create(self._decomp_view.TITLE) def start_decompilation(self, ea: int) -> None: self._current_decomp = None self._current_summary = None self._current_comments = None + self._current_tokenised = None + self._baseline = None self._current_func_vaddr = ea self.run_dialog() @@ -101,6 +115,7 @@ def _dispatch_task(self, ea: int) -> None: on_decomp=lambda response: self._on_decomp_complete(ea, response), on_summary=lambda response: self._on_summary_complete(ea, response), on_comments=lambda response: self._on_comments_complete(ea, response), + on_tokenised=lambda response: self._on_tokenised_complete(ea, response), ) def _on_decomp_complete( @@ -160,80 +175,136 @@ def _on_comments_complete( self._current_comments = response.data self._rerender() - def _rerender(self) -> None: - if self._decomp_view is None or self._current_decomp is None: + def _on_tokenised_complete( + self, ea: int, response: GenericApiReturn[TokenisedData] + ) -> None: + if ea != self._current_func_vaddr: return - rendered = render_view( - decomp=self._current_decomp, - summary=self._current_summary, - comments=self._current_comments, - ) - self._decomp_view.update_view_content(rendered) - - def _on_pane_closed(self) -> None: - self._decomp_view = None - self.disable_function_tracking() - self.log.info("AI Decomp view closed, reference cleared.") - + if not response.success or response.data is None: + return + self._current_tokenised = response.data -def render_view( - decomp: DecompilationData, - summary: SummaryData | None, - comments: CommentsData | None, -) -> str: - code: str = decomp.decompilation or "" + def refresh_current(self) -> None: + ea = self._current_func_vaddr + if ea is None: + return + self.ai_decomp_service.invalidate_ea(ea) + self.start_decompilation(ea) - header_parts: list[str] = [] - if summary is not None and summary.ai_summary: - header_parts.append(_format_summary_as_comment(summary.ai_summary)) + def request_rename(self, display_line: int, word: str) -> None: + ea = self._current_func_vaddr + if ea is None or self._baseline is None or self._current_tokenised is None: + return + if not (0 <= display_line < len(self._baseline.display_is_code)): + return + if not self._baseline.display_is_code[display_line]: + return - body = code - if comments is not None and comments.inline_comments: - body = _inject_inline_comments(code, comments.inline_comments) + source_line = self._baseline.display_source[display_line] + if source_line is None: + return + source_index = source_line - 1 + code_line = self._baseline.code_lines[source_index] + ident_index = index_of_identifier(code_line, word) + if ident_index < 0: + return - if header_parts: - return "\n".join(header_parts + [body]) - return body + resolved = resolve_token(self._current_tokenised, source_index, ident_index, word) + if resolved is None: + self.show_info_dialog( + message=f"'{word}' is not a renameable variable or type." + ) + return + token, _category = resolved + new_name = ida_kernwin.ask_str(word, 0, f"Rename '{word}'") + if not new_name or new_name == word: + return -def _inject_inline_comments(code: str, comments: list[InlineComment]) -> str: - lines = code.split("\n") - for c in sorted(comments, key=lambda x: x.line, reverse=True): - idx = c.line - 1 - if idx < 0 or idx >= len(lines): - continue - target = lines[idx] - indent = target[: len(target) - len(target.lstrip())] - lines.insert(idx, f"{indent}// {c.comment}") - return "\n".join(lines) + self.ai_decomp_service.apply_overrides( + ea=ea, + overrides={token: new_name}, + on_decomp=lambda response: self._on_decomp_complete(ea, response), + on_tokenised=lambda response: self._on_tokenised_complete(ea, response), + ) + def request_edit_comment(self, display_line: int) -> None: + ea = self._current_func_vaddr + if ea is None or self._baseline is None: + return + source_line = self._source_line_for(display_line) + if source_line is None: + return -def _format_summary_as_comment(summary: str) -> str: - prefix = " * " - max_comment_width: int = 100 - content_width: int = max_comment_width - len(prefix) + existing = self._baseline.comment_by_source.get(source_line, "") + text = ida_kernwin.ask_text(0, existing, f"Comment for line {source_line}") + if text is None: + return + text = text.strip() + if not text: + if source_line in self._baseline.comment_by_source: + self.ai_decomp_service.remove_comment( + ea=ea, + line=source_line, + on_result=lambda response: self._on_comments_updated(ea, response), + ) + return - lines: list[str] = ["/*"] + self.ai_decomp_service.set_comment( + ea=ea, + line=source_line, + comment=text, + on_result=lambda response: self._on_comments_updated(ea, response), + ) - for paragraph in summary.split("\n"): - if not paragraph.strip(): - lines.append(" *") - continue + def request_remove_comment(self, display_line: int) -> None: + ea = self._current_func_vaddr + if ea is None or self._baseline is None: + return + source_line = self._source_line_for(display_line) + if source_line is None: + return + if source_line not in self._baseline.comment_by_source: + self.show_info_dialog(message="No comment on this line.") + return + self.ai_decomp_service.remove_comment( + ea=ea, + line=source_line, + on_result=lambda response: self._on_comments_updated(ea, response), + ) - words: list[str] = paragraph.split() - current_line: str = "" + def _source_line_for(self, display_line: int) -> int | None: + if self._baseline is None: + return None + if not (0 <= display_line < len(self._baseline.display_source)): + return None + return self._baseline.display_source[display_line] - for word in words: - if not current_line: - current_line = word - elif len(current_line) + 1 + len(word) <= content_width: - current_line += " " + word - else: - lines.append(prefix + current_line) - current_line = word + def _on_comments_updated( + self, ea: int, response: GenericApiReturn[CommentsData] + ) -> None: + if ea != self._current_func_vaddr: + return + if not response.success: + if response.error_message: + self.show_error_dialog(message=response.error_message) + return + if response.data is None: + return + self._current_comments = response.data + self._rerender() - if current_line: - lines.append(prefix + current_line) + def _rerender(self) -> None: + if self._decomp_view is None or self._current_decomp is None: + return + rendered, self._baseline = render_view_with_map( + decomp=self._current_decomp, + summary=self._current_summary, + comments=self._current_comments, + ) + self._decomp_view.update_view_content(rendered) - lines.append(" */") - return "\n".join(lines) + def _on_pane_closed(self) -> None: + self._decomp_view = None + self.disable_function_tracking() + self.log.info("AI Decomp view closed, reference cleared.") diff --git a/reai_toolkit/app/coordinators/ai_decomp_render.py b/reai_toolkit/app/coordinators/ai_decomp_render.py new file mode 100644 index 00000000..fbf60c6b --- /dev/null +++ b/reai_toolkit/app/coordinators/ai_decomp_render.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping + from revengai.models.comments_data import CommentsData + from revengai.models.decompilation_data import DecompilationData + from revengai.models.summary_data import SummaryData + from revengai.models.tokenised_data import TokenisedData + + +_IDENT_RE = re.compile(r"[A-Za-z_]\w*") +_VARIABLE_CATEGORIES = ( + "unmatched_vars", + "unmatched_global_vars", + "unmatched_external_vars", +) +_TYPE_CATEGORIES = ("unmatched_custom_types", "unmatched_enums") + + +@dataclass +class RenderModel: + summary_line_count: int + code_lines: list[str] + comment_by_source: dict[int, str] + display_source: list[Optional[int]] + display_is_code: list[bool] + + +def render_view( + decomp: "DecompilationData", + summary: "Optional[SummaryData]", + comments: "Optional[CommentsData]", +) -> str: + text, _ = render_view_with_map(decomp, summary, comments) + return text + + +def render_view_with_map( + decomp: "DecompilationData", + summary: "Optional[SummaryData]", + comments: "Optional[CommentsData]", +) -> tuple[str, RenderModel]: + code: str = decomp.decompilation or "" + code_lines: list[str] = code.split("\n") + + summary_block: str | None = None + if summary is not None and summary.ai_summary: + summary_block = _format_summary_as_comment(summary.ai_summary) + summary_line_count = len(summary_block.split("\n")) if summary_block is not None else 0 + + inline = comments.inline_comments if comments is not None else None + comment_by_source: dict[int, str] = {} + if inline: + for c in inline: + if 1 <= c.line <= len(code_lines): + comment_by_source[c.line] = c.comment + + display_lines: list[str] = [] + display_source: list[Optional[int]] = [] + display_is_code: list[bool] = [] + + if summary_block is not None: + for sline in summary_block.split("\n"): + display_lines.append(sline) + display_source.append(None) + display_is_code.append(False) + + for idx, code_line in enumerate(code_lines): + source_line = idx + 1 + comment = comment_by_source.get(source_line) + if comment is not None: + indent = code_line[: len(code_line) - len(code_line.lstrip())] + for part in comment.split("\n"): + display_lines.append(f"{indent}// {part}") + display_source.append(source_line) + display_is_code.append(False) + display_lines.append(code_line) + display_source.append(source_line) + display_is_code.append(True) + + text = "\n".join(display_lines) + return text, RenderModel( + summary_line_count=summary_line_count, + code_lines=code_lines, + comment_by_source=comment_by_source, + display_source=display_source, + display_is_code=display_is_code, + ) + + +def _format_summary_as_comment(summary: str) -> str: + prefix = " * " + max_comment_width: int = 100 + content_width: int = max_comment_width - len(prefix) + + lines: list[str] = ["/*"] + + for paragraph in summary.split("\n"): + if not paragraph.strip(): + lines.append(" *") + continue + + words: list[str] = paragraph.split() + current_line: str = "" + + for word in words: + if not current_line: + current_line = word + elif len(current_line) + 1 + len(word) <= content_width: + current_line += " " + word + else: + lines.append(prefix + current_line) + current_line = word + + if current_line: + lines.append(prefix + current_line) + + lines.append(" */") + return "\n".join(lines) + + +def index_of_identifier(line: str, word: str) -> int: + idents = _IDENT_RE.findall(line) + return idents.index(word) if word in idents else -1 + + +def resolve_token( + tokenised: "TokenisedData", + source_index: int, + ident_index: int, + old_ident: str, +) -> Optional[tuple[str, str]]: + mapping = tokenised.function_mapping + if mapping is None: + return None + + tok_lines = (tokenised.tokenised_decompilation or "").split("\n") + if 0 <= source_index < len(tok_lines): + tok_idents = _IDENT_RE.findall(tok_lines[source_index]) + if 0 <= ident_index < len(tok_idents): + candidate = tok_idents[ident_index] + cat, eff = _category_of_token(mapping, candidate) + if cat is not None and eff == old_ident: + return candidate, cat + + matches = [ + (token, cat) + for token, rv, cat in _iter_category_tokens(mapping) + if _effective_value(mapping, token, rv) == old_ident + ] + if len(matches) == 1: + return matches[0] + return None + + +def _iter_category_tokens(mapping: "AIDecompFunctionMapping"): + for cat in _VARIABLE_CATEGORIES: + for token, rv in (getattr(mapping, cat, None) or {}).items(): + yield token, rv, "variable" + for cat in _TYPE_CATEGORIES: + for token, rv in (getattr(mapping, cat, None) or {}).items(): + yield token, rv, "type" + + +def _category_of_token( + mapping: "AIDecompFunctionMapping", token: str +) -> tuple[Optional[str], Optional[str]]: + for t, rv, cat in _iter_category_tokens(mapping): + if t == token: + return cat, _effective_value(mapping, t, rv) + return None, None + + +def _effective_value(mapping: "AIDecompFunctionMapping", token: str, replacement) -> str: + overrides = mapping.user_override_mappings or {} + return overrides.get(token, replacement.value) diff --git a/reai_toolkit/app/core/logging_setup.py b/reai_toolkit/app/core/logging_setup.py new file mode 100644 index 00000000..6bc0d55e --- /dev/null +++ b/reai_toolkit/app/core/logging_setup.py @@ -0,0 +1,12 @@ +import os +import sys + +from loguru import logger + + +def configure_logging() -> None: + level = os.environ.get("REAI_LOG_LEVEL") + if not level: + level = "DEBUG" if os.environ.get("REAI_DEBUG") == "1" else "INFO" + logger.remove() + logger.add(sys.stderr, level=level.upper()) diff --git a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py index 5b69b74e..0a871b75 100644 --- a/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py +++ b/reai_toolkit/app/services/ai_decomp/ai_decomp_service.py @@ -10,8 +10,11 @@ ) from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData +from revengai.models.patch_comment_body import PatchCommentBody from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus +from revengai.models.tokenised_data import TokenisedData +from revengai.models.upsert_overrides_input_body import UpsertOverridesInputBody from revengai.models.workflow_progress import WorkflowProgress from reai_toolkit.app.core.netstore_service import SimpleNetStore @@ -30,6 +33,7 @@ def __init__(self, netstore_service: SimpleNetStore, sdk_config: Configuration) self._decomp_cache: dict[int, DecompilationData] = {} self._summary_cache: dict[int, SummaryData] = {} self._comments_cache: dict[int, CommentsData] = {} + self._tokenised_cache: dict[int, TokenisedData] = {} self._inflight: dict[int, threading.Event] = {} self._inflight_lock = threading.Lock() @@ -49,6 +53,7 @@ def stop_worker(self) -> None: self._decomp_cache.clear() self._summary_cache.clear() self._comments_cache.clear() + self._tokenised_cache.clear() def peek_decomp(self, ea: int) -> DecompilationData | None: function_id = self._get_function_id(start_ea=ea) @@ -56,12 +61,31 @@ def peek_decomp(self, ea: int) -> DecompilationData | None: return None return self._decomp_cache.get(function_id) + def function_id_for(self, ea: int) -> int | None: + return self._get_function_id(start_ea=ea) + + def invalidate(self, function_id: int) -> None: + self._decomp_cache.pop(function_id, None) + self._summary_cache.pop(function_id, None) + self._comments_cache.pop(function_id, None) + self._tokenised_cache.pop(function_id, None) + with self._inflight_lock: + evt = self._inflight.pop(function_id, None) + if evt is not None: + evt.set() + + def invalidate_ea(self, ea: int) -> None: + function_id = self._get_function_id(start_ea=ea) + if function_id is not None: + self.invalidate(function_id) + def start_ai_decomp_task( self, ea: int, on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], on_summary: Callable[[GenericApiReturn[SummaryData]], None], on_comments: Callable[[GenericApiReturn[CommentsData]], None], + on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None] | None = None, ) -> None: function_id = self._get_function_id(start_ea=ea) if function_id is None: @@ -76,12 +100,201 @@ def start_ai_decomp_task( worker = threading.Thread( target=self._run_task, - args=(function_id, stop_event, on_decomp, on_summary, on_comments), + args=( + function_id, + stop_event, + on_decomp, + on_summary, + on_comments, + on_tokenised, + ), name=f"reai-aidecomp-{function_id}", daemon=True, ) worker.start() + def apply_overrides( + self, + ea: int, + overrides: dict[str, str], + on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], + on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None], + ) -> None: + function_id = self._get_function_id(start_ea=ea) + if function_id is None: + on_decomp( + GenericApiReturn[DecompilationData]( + success=False, error_message="Function is not part of the analysis." + ) + ) + return + self._spawn( + self._run_apply_overrides, + f"reai-aidecomp-overrides-{function_id}", + (function_id, overrides, on_decomp, on_tokenised), + ) + + def set_comment( + self, + ea: int, + line: int, + comment: str, + on_result: Callable[[GenericApiReturn[CommentsData]], None], + ) -> None: + function_id = self._get_function_id(start_ea=ea) + if function_id is None: + on_result( + GenericApiReturn[CommentsData]( + success=False, error_message="Function is not part of the analysis." + ) + ) + return + self._spawn( + self._run_set_comment, + f"reai-aidecomp-comment-{function_id}", + (function_id, line, comment, on_result), + ) + + def remove_comment( + self, + ea: int, + line: int, + on_result: Callable[[GenericApiReturn[CommentsData]], None], + ) -> None: + function_id = self._get_function_id(start_ea=ea) + if function_id is None: + on_result( + GenericApiReturn[CommentsData]( + success=False, error_message="Function is not part of the analysis." + ) + ) + return + self._spawn( + self._run_remove_comment, + f"reai-aidecomp-comment-del-{function_id}", + (function_id, line, on_result), + ) + + @staticmethod + def _spawn(target: Callable[..., Any], name: str, args: tuple) -> None: + threading.Thread(target=target, args=args, name=name, daemon=True).start() + + def _run_apply_overrides( + self, + function_id: int, + overrides: dict[str, str], + on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], + on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None], + ) -> None: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + FunctionsAIDecompilationApi(api_client).upsert_ai_decompilation_overrides( + function_id=function_id, + upsert_overrides_input_body=UpsertOverridesInputBody( + overrides=overrides + ), + ) + except ApiException as e: + on_decomp( + GenericApiReturn[DecompilationData]( + success=False, error_message=_format_api_error(e) + ) + ) + return + except Exception as e: + on_decomp( + GenericApiReturn[DecompilationData]( + success=False, error_message=f"Unexpected error applying overrides: {e}" + ) + ) + return + + self.invalidate(function_id) + + decomp, derr = self._fetch_decompilation(function_id) + if decomp is not None and decomp.decompilation: + self._decomp_cache[function_id] = decomp + on_decomp(GenericApiReturn[DecompilationData](success=True, data=decomp)) + else: + on_decomp( + GenericApiReturn[DecompilationData]( + success=False, + error_message=derr or "AI decompilation returned no content.", + ) + ) + + tokenised, _terr = self._fetch_tokenised(function_id) + if ( + tokenised is not None + and str(tokenised.status) == TaskStatus.COMPLETED.value + and tokenised.function_mapping is not None + ): + self._tokenised_cache[function_id] = tokenised + on_tokenised(GenericApiReturn[TokenisedData](success=True, data=tokenised)) + + def _run_set_comment( + self, + function_id: int, + line: int, + comment: str, + on_result: Callable[[GenericApiReturn[CommentsData]], None], + ) -> None: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + data = FunctionsAIDecompilationApi( + api_client + ).patch_ai_decompilation_inline_comment( + function_id=function_id, + patch_comment_body=PatchCommentBody(comment=comment, line=line), + ) + except ApiException as e: + on_result( + GenericApiReturn[CommentsData]( + success=False, error_message=_format_api_error(e) + ) + ) + return + except Exception as e: + on_result( + GenericApiReturn[CommentsData]( + success=False, error_message=f"Unexpected error updating comment: {e}" + ) + ) + return + self._comments_cache[function_id] = data + on_result(GenericApiReturn[CommentsData](success=True, data=data)) + + def _run_remove_comment( + self, + function_id: int, + line: int, + on_result: Callable[[GenericApiReturn[CommentsData]], None], + ) -> None: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + data = FunctionsAIDecompilationApi( + api_client + ).delete_ai_decompilation_inline_comment( + function_id=function_id, + line=line, + ) + except ApiException as e: + on_result( + GenericApiReturn[CommentsData]( + success=False, error_message=_format_api_error(e) + ) + ) + return + except Exception as e: + on_result( + GenericApiReturn[CommentsData]( + success=False, error_message=f"Unexpected error removing comment: {e}" + ) + ) + return + self._comments_cache[function_id] = data + on_result(GenericApiReturn[CommentsData](success=True, data=data)) + def _run_task( self, function_id: int, @@ -89,6 +302,7 @@ def _run_task( on_decomp: Callable[[GenericApiReturn[DecompilationData]], None], on_summary: Callable[[GenericApiReturn[SummaryData]], None], on_comments: Callable[[GenericApiReturn[CommentsData]], None], + on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None] | None = None, ) -> None: try: if stop_event.is_set(): @@ -98,11 +312,14 @@ def _run_task( return self._run_summary_phase(function_id, stop_event, on_summary) self._run_comments_phase(function_id, stop_event, on_comments) + if on_tokenised is not None: + self._run_tokenised_phase(function_id, stop_event, on_tokenised) except Exception as e: logger.error(f"RevEng.AI: AI decompilation task crashed for {function_id}: {e}") finally: with self._inflight_lock: - self._inflight.pop(function_id, None) + if self._inflight.get(function_id) is stop_event: + self._inflight.pop(function_id, None) def _get_function_id(self, start_ea: int) -> int | None: function_map: FunctionMapping | None = self.netstore_service.get_function_mapping() @@ -501,6 +718,63 @@ def _run_comments_phase( GenericApiReturn[CommentsData](success=True, data=final), ) + def _fetch_tokenised( + self, function_id: int + ) -> tuple[TokenisedData | None, str | None]: + try: + with self.yield_api_client(sdk_config=self.sdk_config) as api_client: + tokenised = FunctionsAIDecompilationApi( + api_client + ).get_ai_decompilation_tokenised(function_id=function_id) + return tokenised, None + except ApiException as e: + return None, _format_api_error(e) + except Exception as e: + return None, f"Unexpected error fetching tokenised data: {e}" + + def _run_tokenised_phase( + self, + function_id: int, + stop_event: threading.Event, + on_tokenised: Callable[[GenericApiReturn[TokenisedData]], None], + ) -> None: + cached = self._tokenised_cache.get(function_id) + if cached is not None: + self._safe_dispatch( + stop_event, + on_tokenised, + GenericApiReturn[TokenisedData](success=True, data=cached), + ) + return + + tokenised, err = self._fetch_tokenised(function_id) + if tokenised is None: + self._safe_dispatch( + stop_event, + on_tokenised, + GenericApiReturn[TokenisedData](success=False, error_message=err), + ) + return + + if ( + str(tokenised.status) == TaskStatus.COMPLETED.value + and tokenised.function_mapping is not None + ): + self._tokenised_cache[function_id] = tokenised + self._safe_dispatch( + stop_event, + on_tokenised, + GenericApiReturn[TokenisedData](success=True, data=tokenised), + ) + else: + self._safe_dispatch( + stop_event, + on_tokenised, + GenericApiReturn[TokenisedData]( + success=False, error_message="Tokenised data not ready." + ), + ) + def _poll_workflow( self, function_id: int, diff --git a/reai_toolkit_entry.py b/reai_toolkit_entry.py index a0bd60b2..4b98efa5 100644 --- a/reai_toolkit_entry.py +++ b/reai_toolkit_entry.py @@ -33,6 +33,7 @@ def _add_vendor_paths(): from loguru import logger from reai_toolkit.app.app import App +from reai_toolkit.app.core.logging_setup import configure_logging from reai_toolkit.app.coordinator import Coordinator from reai_toolkit.app.factory import DialogFactory from reai_toolkit.hooks.artifacts import ArtifactChangeHooks @@ -70,9 +71,9 @@ def __init__(self): self._setup_handled = False def init(self): + configure_logging() self.app = App(__IDA_VERSION__, __PLUGIN_VERSION__) self.factory = DialogFactory(self.app) - print(Coordinator.__abstractmethods__) self.coordinator = Coordinator(app=self.app, factory=self.factory, log=logger) # Top level menu hooks diff --git a/tests/idalib/ai_decomp/test_smoke.py b/tests/idalib/ai_decomp/test_smoke.py index d05b7a8b..37610e87 100644 --- a/tests/idalib/ai_decomp/test_smoke.py +++ b/tests/idalib/ai_decomp/test_smoke.py @@ -6,8 +6,10 @@ import idautils from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData +from revengai.models.inline_comment import InlineComment from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus +from revengai.models.tokenised_data import TokenisedData from reai_toolkit.app.services.ai_decomp import ai_decomp_service as svc_mod from reai_toolkit.app.services.ai_decomp.ai_decomp_service import AiDecompService @@ -54,3 +56,69 @@ def test_service_runs_to_completion_under_idalib(loaded_binary, mocker): assert result.success is True assert result.data.decompilation == "int sub() { return 0; }" api_inst.get_ai_decompilation.assert_called_once_with(function_id=7) + + +def _wait_mock(mock, timeout=10.0): + deadline = time.monotonic() + timeout + while not mock.called and time.monotonic() < deadline: + time.sleep(0.01) + assert mock.called, "callback did not fire under idalib" + + +def test_mutations_and_refresh_under_idalib(loaded_binary, mocker): + ea = next(iter(idautils.Functions()), None) + assert ea is not None + + mocker.patch.object(svc_mod.AiDecompService, "yield_api_client") + api_inst = MagicMock() + mocker.patch.object(svc_mod, "FunctionsAIDecompilationApi", return_value=api_inst) + api_inst.get_ai_decompilation.return_value = DecompilationData.model_construct( + status=TaskStatus.COMPLETED.value, decompilation="int sub(int v1) { return v1; }" + ) + api_inst.get_ai_decompilation_tokenised.return_value = TokenisedData.model_construct( + status=TaskStatus.COMPLETED.value, + tokenised_decompilation="int @@F@@(int @@V@@) { return @@V@@; }", + predicted_function_name="sub", + function_mapping=MagicMock(), + ) + api_inst.patch_ai_decompilation_inline_comment.return_value = ( + CommentsData.model_construct( + inline_comments=[InlineComment.model_construct(comment="note", line=1)], + task_status=TaskStatus.COMPLETED.value, + ) + ) + api_inst.delete_ai_decompilation_inline_comment.return_value = ( + CommentsData.model_construct( + inline_comments=[], task_status=TaskStatus.COMPLETED.value + ) + ) + + netstore = MagicMock() + netstore.get_function_mapping.return_value.inverse_function_map = {str(ea): 7} + service = AiDecompService(netstore_service=netstore, sdk_config=MagicMock()) + + on_decomp, on_tok = MagicMock(), MagicMock() + service.apply_overrides( + ea=ea, overrides={"@@V@@": "count"}, on_decomp=on_decomp, on_tokenised=on_tok + ) + _wait_mock(on_decomp) + _wait_mock(on_tok) + api_inst.upsert_ai_decompilation_overrides.assert_called_once() + assert on_decomp.call_args[0][0].success is True + assert service._tokenised_cache[7] is not None + + on_set = MagicMock() + service.set_comment(ea=ea, line=1, comment="note", on_result=on_set) + _wait_mock(on_set) + assert on_set.call_args[0][0].success is True + assert service._comments_cache[7] is not None + + on_del = MagicMock() + service.remove_comment(ea=ea, line=1, on_result=on_del) + _wait_mock(on_del) + assert on_del.call_args[0][0].success is True + + service.invalidate_ea(ea) + assert 7 not in service._decomp_cache + assert 7 not in service._tokenised_cache + assert 7 not in service._comments_cache diff --git a/tests/ui/ai_decomp_ui_runner.py b/tests/ui/ai_decomp_ui_runner.py new file mode 100644 index 00000000..6cc5b437 --- /dev/null +++ b/tests/ui/ai_decomp_ui_runner.py @@ -0,0 +1,218 @@ +import json +import os +import sys +import traceback + + +def _add_paths() -> None: + root = os.environ["REAI_UI_ROOT"] + pyver = f"python{sys.version_info.major}.{sys.version_info.minor}" + suffixes = [ + "vendor", + "vendor/site-packages", + "vendor/Lib/site-packages", + f"vendor/lib/{pyver}/site-packages", + f"vendor/{pyver}/site-packages", + ] + for base in (root, os.path.join(root, "reai_toolkit")): + for suffix in suffixes: + path = os.path.join(base, suffix) + if os.path.isdir(path) and path not in sys.path: + sys.path.insert(0, path) + if root not in sys.path: + sys.path.insert(0, root) + + +CODE = "int f(int a1) {\n int v5 = a1;\n return v5;\n}" +TOK = "int @@F@@(int @@A@@) {\n int @@V@@ = @@A@@;\n return @@V@@;\n}" + + +def _run(report: dict) -> None: + from types import SimpleNamespace + from unittest.mock import MagicMock + + import ida_kernwin + import idautils + + from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping + from revengai.models.comments_data import CommentsData + from revengai.models.decompilation_data import DecompilationData + from revengai.models.inline_comment import InlineComment + from revengai.models.replacement_value import ReplacementValue + from revengai.models.tokenised_data import TokenisedData + + from reai_toolkit.app.components.tabs.ai_decomp_tab import AIDecompView + from reai_toolkit.app.coordinators.ai_decomp_coordinator import AiDecompCoordinator + from reai_toolkit.app.core.qt_compat import QtWidgets + from reai_toolkit.app.core.shared_schema import GenericApiReturn + + answers = {"str": "count", "text": "hello"} + ida_kernwin.ask_str = lambda default, hist, prompt: answers["str"] + ida_kernwin.ask_text = lambda maxsz, default, prompt: answers["text"] + + ea = next(iter(idautils.Functions()), 0x1000) + + decomp = DecompilationData.model_construct(status="COMPLETED", decompilation=CODE) + mapping = AIDecompFunctionMapping.model_construct( + unmatched_vars={ + "@@A@@": ReplacementValue.model_construct(value="a1"), + "@@V@@": ReplacementValue.model_construct(value="v5"), + }, + user_override_mappings={}, + ) + tokenised = TokenisedData.model_construct( + status="COMPLETED", + tokenised_decompilation=TOK, + predicted_function_name="f", + function_mapping=mapping, + ) + + service = MagicMock() + service.peek_decomp.return_value = None + factory = SimpleNamespace( + ai_decomp=lambda on_closed: AIDecompView(on_closed=on_closed) + ) + coord = AiDecompCoordinator( + app=MagicMock(), factory=factory, log=MagicMock(), ai_decomp_service=service + ) + infos: list = [] + coord.show_info_dialog = lambda **kw: infos.append(kw) + coord.show_error_dialog = lambda **kw: infos.append(kw) + + qapp = QtWidgets.QApplication.instance() + + def pump() -> None: + if qapp is not None: + for _ in range(5): + qapp.processEvents() + + coord.run_dialog() + pump() + view = coord._decomp_view + if view is None or view._editor is None: + report["errors"].append("AIDecompView editor was not created") + return + report["view_created"] = True + report["editor_read_only"] = view._editor.isReadOnly() is True + + def seed_plain() -> None: + service.reset_mock() + service.peek_decomp.return_value = None + coord._current_func_vaddr = ea + coord._current_summary = None + coord._current_comments = None + coord._current_decomp = decomp + coord._current_tokenised = tokenised + coord._rerender() + pump() + + def seed_with_comment() -> None: + service.reset_mock() + service.peek_decomp.return_value = None + coord._current_func_vaddr = ea + coord._current_summary = None + coord._current_decomp = decomp + coord._current_tokenised = tokenised + coord._on_comments_complete( + ea, + GenericApiReturn( + success=True, + data=CommentsData.model_construct( + inline_comments=[InlineComment.model_construct(comment="hola", line=2)], + task_status="COMPLETED", + ), + ), + ) + pump() + + def code_line_row(needle: str) -> int: + lines = view._editor.toPlainText().split("\n") + return next(i for i, s in enumerate(lines) if needle in s) + + seed_plain() + report["render_shows_code"] = CODE in view._editor.toPlainText() + + seed_plain() + answers["str"] = "count" + view._editor.renameRequested.emit(code_line_row("int v5"), "v5") + pump() + report["rename_double_click_overrides"] = service.apply_overrides.called + if service.apply_overrides.called: + report["rename_overrides_correct"] = ( + service.apply_overrides.call_args.kwargs.get("overrides") == {"@@V@@": "count"} + ) + + seed_plain() + infos.clear() + view._editor.renameRequested.emit(0, "int") + pump() + report["rename_non_token_info"] = not service.apply_overrides.called and len(infos) >= 1 + + seed_plain() + answers["text"] = "hello" + view._editor.commentEditRequested.emit(code_line_row("int v5")) + pump() + report["comment_add_sets"] = service.set_comment.called + if service.set_comment.called: + kw = service.set_comment.call_args.kwargs + report["comment_add_args_correct"] = kw.get("line") == 2 and kw.get("comment") == "hello" + + seed_with_comment() + answers["text"] = "" + view._editor.commentEditRequested.emit(code_line_row("int v5")) + pump() + report["comment_edit_empty_removes"] = ( + service.remove_comment.called + and service.remove_comment.call_args.kwargs.get("line") == 2 + ) + + seed_with_comment() + view._editor.commentRemoveRequested.emit(code_line_row("int v5")) + pump() + report["comment_remove_deletes"] = service.remove_comment.called + if service.remove_comment.called: + report["comment_remove_args_correct"] = ( + service.remove_comment.call_args.kwargs.get("line") == 2 + ) + + service.reset_mock() + service.peek_decomp.return_value = None + view._refresh_btn.click() + pump() + report["refresh_button_invalidates"] = service.invalidate_ea.called + + report["ok"] = not report["errors"] + + +def main() -> None: + import ida_auto + import ida_pro + + report = { + "ok": False, + "errors": [], + "view_created": False, + "editor_read_only": False, + "render_shows_code": False, + "rename_double_click_overrides": False, + "rename_overrides_correct": False, + "rename_non_token_info": False, + "comment_add_sets": False, + "comment_add_args_correct": False, + "comment_edit_empty_removes": False, + "comment_remove_deletes": False, + "comment_remove_args_correct": False, + "refresh_button_invalidates": False, + } + ida_auto.auto_wait() + try: + _add_paths() + _run(report) + except Exception: + report["errors"].append(traceback.format_exc()) + with open(os.environ["REAI_UI_REPORT"], "w") as fh: + json.dump(report, fh) + ida_pro.qexit(0) + + +main() diff --git a/tests/ui/test_gui_ai_decomp.py b/tests/ui/test_gui_ai_decomp.py new file mode 100644 index 00000000..c6084955 --- /dev/null +++ b/tests/ui/test_gui_ai_decomp.py @@ -0,0 +1,101 @@ +import json +import os +import shutil +import subprocess +import sys + +import pytest + +pytestmark = pytest.mark.ida_ui + +ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +RUNNER = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ai_decomp_ui_runner.py") +HELLO_ELF = os.path.join(ROOT, "tests", "fixtures", "hello.elf") + +CHECKS = [ + "view_created", + "editor_read_only", + "render_shows_code", + "rename_double_click_overrides", + "rename_overrides_correct", + "rename_non_token_info", + "comment_add_sets", + "comment_add_args_correct", + "comment_edit_empty_removes", + "comment_remove_deletes", + "comment_remove_args_correct", + "refresh_button_invalidates", +] + + +def _ida_gui_binary() -> str | None: + path = os.path.join(os.environ.get("IDADIR", ""), "ida") + return path if os.path.isfile(path) else None + + +def _ida_gui_running() -> bool: + for name in ("ida", "ida64", "idat"): + try: + if subprocess.run(["pgrep", "-x", name], capture_output=True).returncode == 0: + return True + except FileNotFoundError: + return False + return False + + +def _headless_env() -> dict[str, str]: + if sys.platform == "linux" and not os.environ.get("DISPLAY") and not os.environ.get("WAYLAND_DISPLAY"): + return {"QT_QPA_PLATFORM": "offscreen"} + return {} + + +@pytest.mark.skipif( + os.environ.get("REAI_UI_TESTS") != "1", + reason="GUI IDA test; set REAI_UI_TESTS=1 to enable", +) +def test_ai_decomp_editing_flow_in_gui(tmp_path): + ida = _ida_gui_binary() + if ida is None: + pytest.skip("IDADIR not set or ida binary missing") + if _ida_gui_running(): + pytest.skip("an IDA instance is already running (license seat busy)") + if not os.path.isfile(HELLO_ELF): + pytest.skip(f"missing fixture {HELLO_ELF}") + + binary = tmp_path / "hello.elf" + shutil.copy(HELLO_ELF, binary) + report_path = tmp_path / "report.json" + log_path = tmp_path / "ida.log" + env = dict( + os.environ, + REAI_UI_REPORT=str(report_path), + REAI_UI_ROOT=ROOT, + **_headless_env(), + ) + + proc = subprocess.run( + [ida, "-A", f"-S{RUNNER}", f"-L{log_path}", str(binary)], + env=env, + capture_output=True, + timeout=300, + ) + + if not report_path.is_file(): + log = log_path.read_text() if log_path.is_file() else "" + stderr = proc.stderr.decode(errors="replace") + stdout = proc.stdout.decode(errors="replace") + if "license" in log.lower() or "license" in stderr.lower(): + pytest.skip("IDA license unavailable") + pytest.fail( + f"IDA produced no report (rc={proc.returncode});" + f"\nstderr tail:\n{stderr[-2000:]}" + f"\nstdout tail:\n{stdout[-2000:]}" + f"\nlog tail:\n{log[-2000:]}" + ) + + report = json.loads(report_path.read_text()) + + assert report["errors"] == [], report["errors"] + assert report["ok"] is True + for check in CHECKS: + assert report.get(check) is True, f"{check} failed: {report}" diff --git a/tests/unit/ai_decomp/test_render.py b/tests/unit/ai_decomp/test_render.py new file mode 100644 index 00000000..38be08e7 --- /dev/null +++ b/tests/unit/ai_decomp/test_render.py @@ -0,0 +1,154 @@ +from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping +from revengai.models.comments_data import CommentsData +from revengai.models.decompilation_data import DecompilationData +from revengai.models.inline_comment import InlineComment +from revengai.models.replacement_value import ReplacementValue +from revengai.models.summary_data import SummaryData +from revengai.models.tokenised_data import TokenisedData + +from reai_toolkit.app.coordinators.ai_decomp_render import ( + index_of_identifier, + render_view, + render_view_with_map, + resolve_token, +) + + +CODE = "int f(int a1) {\n int v5 = a1;\n return v5;\n}" +TOK = "int @@F1@@(int @@V_a1@@) {\n int @@V_v5@@ = @@V_a1@@;\n return @@V_v5@@;\n}" + + +def _dd(code=CODE): + return DecompilationData.model_construct(status="COMPLETED", decompilation=code) + + +def _summary(text): + return SummaryData.model_construct(ai_summary=text, summary=text, task_status="COMPLETED") + + +def _comments(pairs): + items = [InlineComment.model_construct(comment=c, line=ln) for ln, c in pairs] + return CommentsData.model_construct(inline_comments=items, task_status="COMPLETED") + + +def _mapping(**cats): + base = dict( + fields={}, + inverse_function_map={}, + inverse_string_map={}, + unmatched_custom_function_pointers={}, + unmatched_custom_types={}, + unmatched_enums={}, + unmatched_external_vars={}, + unmatched_functions={}, + unmatched_global_vars={}, + unmatched_go_to_labels={}, + unmatched_strings={}, + unmatched_variadic_lists={}, + unmatched_vars={}, + user_override_mappings={}, + ) + base.update(cats) + return AIDecompFunctionMapping.model_construct(**base) + + +def _rv(value): + return ReplacementValue.model_construct(value=value) + + +def _tokenised(tok=TOK, mapping=None): + return TokenisedData.model_construct( + status="COMPLETED", + tokenised_decompilation=tok, + predicted_function_name="f", + function_mapping=mapping if mapping is not None else _mapping(), + ) + + +def test_render_matches_legacy_and_builds_model(): + comments = _comments([(2, "local copy"), (99, "out of range")]) + text, model = render_view_with_map(_dd(), _summary("Adds one."), comments) + + assert render_view(_dd(), _summary("Adds one."), comments) == text + assert text.splitlines()[0] == "/*" + assert " // local copy" in text + assert model.summary_line_count == 3 + assert model.code_lines == CODE.split("\n") + assert model.comment_by_source == {2: "local copy"} + + +def test_display_maps_align_lines_to_source(): + text, model = render_view_with_map(_dd(), _summary("S."), _comments([(2, "note")])) + lines = text.split("\n") + + assert len(model.display_source) == len(lines) + assert len(model.display_is_code) == len(lines) + + for i in range(model.summary_line_count): + assert model.display_source[i] is None + assert model.display_is_code[i] is False + + comment_rows = [i for i, s in enumerate(lines) if s.strip().startswith("//")] + code_rows = [i for i, c in enumerate(model.display_is_code) if c] + + assert len(code_rows) == len(CODE.split("\n")) + for row in comment_rows: + assert model.display_is_code[row] is False + assert model.display_source[row] == 2 + + v5_row = next(i for i, s in enumerate(lines) if "int v5" in s) + assert model.display_is_code[v5_row] is True + assert model.display_source[v5_row] == 2 + + +def test_no_summary_no_comments_is_raw_code(): + text, model = render_view_with_map(_dd(), None, None) + assert text == CODE + assert model.summary_line_count == 0 + assert model.comment_by_source == {} + assert all(model.display_is_code) + assert model.display_source == [1, 2, 3, 4] + + +def test_multiline_comment_renders_as_two_comment_lines(): + text, model = render_view_with_map(_dd(), None, _comments([(2, "line one\nline two")])) + lines = text.split("\n") + assert " // line one" in lines + assert " // line two" in lines + assert model.comment_by_source == {2: "line one\nline two"} + + +def test_index_of_identifier(): + assert index_of_identifier(" int v5 = a1;", "v5") == 1 + assert index_of_identifier(" int v5 = a1;", "a1") == 2 + assert index_of_identifier(" int v5 = a1;", "missing") == -1 + + +def test_resolve_token_positional_variable(): + tokd = _tokenised(mapping=_mapping(unmatched_vars={"@@V_v5@@": _rv("v5"), "@@V_a1@@": _rv("a1")})) + assert resolve_token(tokd, 1, 1, "v5") == ("@@V_v5@@", "variable") + + +def test_resolve_token_type_category(): + tok = "@@T_S@@ *x = 0;" + tokd = _tokenised(tok=tok, mapping=_mapping(unmatched_custom_types={"@@T_S@@": _rv("Foo")})) + assert resolve_token(tokd, 0, 0, "Foo") == ("@@T_S@@", "type") + + +def test_resolve_token_honours_user_override(): + mapping = _mapping( + unmatched_vars={"@@V_v5@@": _rv("v5")}, + user_override_mappings={"@@V_v5@@": "tmp"}, + ) + tokd = _tokenised(mapping=mapping) + assert resolve_token(tokd, 1, 1, "tmp") == ("@@V_v5@@", "variable") + + +def test_resolve_token_unknown_returns_none(): + tokd = _tokenised(mapping=_mapping(unmatched_vars={"@@V_v5@@": _rv("v5")})) + assert resolve_token(tokd, 1, 1, "not_a_var") is None + + +def test_resolve_token_value_fallback_when_line_unaligned(): + tokd = _tokenised(tok="", mapping=_mapping(unmatched_vars={"@@V_v5@@": _rv("v5")})) + assert resolve_token(tokd, 1, 1, "v5") == ("@@V_v5@@", "variable") diff --git a/tests/unit/ai_decomp/test_sdk_schemas.py b/tests/unit/ai_decomp/test_sdk_schemas.py index e7ebce07..e3083f95 100644 --- a/tests/unit/ai_decomp/test_sdk_schemas.py +++ b/tests/unit/ai_decomp/test_sdk_schemas.py @@ -2,9 +2,16 @@ Schema-shape assertions: fail loudly when the revengai SDK shape drifts. """ +from revengai.models.ai_decomp_function_mapping import AIDecompFunctionMapping +from revengai.models.comments_data import CommentsData from revengai.models.create_ai_decomp_output_body import CreateAIDecompOutputBody +from revengai.models.inline_comment import InlineComment +from revengai.models.patch_comment_body import PatchCommentBody +from revengai.models.replacement_value import ReplacementValue from revengai.models.task_status import TaskStatus from revengai.models.tokenised_data import TokenisedData +from revengai.models.upsert_overrides_data import UpsertOverridesData +from revengai.models.upsert_overrides_input_body import UpsertOverridesInputBody from revengai.models.workflow_progress import WorkflowProgress @@ -37,3 +44,42 @@ def test_workflow_progress_accepts_uppercase_status_values(): WorkflowProgress.model_construct( status=status, step="x", step_index=0, steps_total=1, messages=[] ) + + +def test_upsert_overrides_input_body_round_trips(): + body = UpsertOverridesInputBody.from_json('{"overrides": {"@@V@@": "buf", "@@T@@": ""}}') + assert body.overrides == {"@@V@@": "buf", "@@T@@": ""} + assert "overrides" in UpsertOverridesInputBody.model_fields + + +def test_upsert_overrides_data_exposes_merged_map(): + data = UpsertOverridesData.from_json('{"user_override_mappings": {"@@V@@": "buf"}}') + assert data.user_override_mappings == {"@@V@@": "buf"} + + +def test_patch_comment_body_round_trips(): + body = PatchCommentBody.from_json('{"comment": "note", "line": 7}') + assert body.comment == "note" + assert body.line == 7 + assert {"comment", "line"} <= set(PatchCommentBody.model_fields.keys()) + + +def test_inline_comment_and_comments_data_shapes(): + assert {"comment", "line"} <= set(InlineComment.model_fields.keys()) + assert {"inline_comments", "task_status"} <= set(CommentsData.model_fields.keys()) + + +def test_replacement_value_carries_value(): + assert ReplacementValue.from_json('{"value": "count"}').value == "count" + + +def test_function_mapping_has_variable_and_type_categories(): + fields = set(AIDecompFunctionMapping.model_fields.keys()) + assert { + "unmatched_vars", + "unmatched_global_vars", + "unmatched_external_vars", + "unmatched_custom_types", + "unmatched_enums", + "user_override_mappings", + } <= fields diff --git a/tests/unit/ai_decomp/test_service.py b/tests/unit/ai_decomp/test_service.py index 3c212763..2e5b8e3d 100644 --- a/tests/unit/ai_decomp/test_service.py +++ b/tests/unit/ai_decomp/test_service.py @@ -6,8 +6,10 @@ from revengai import ApiException from revengai.models.comments_data import CommentsData from revengai.models.decompilation_data import DecompilationData +from revengai.models.inline_comment import InlineComment from revengai.models.summary_data import SummaryData from revengai.models.task_status import TaskStatus +from revengai.models.tokenised_data import TokenisedData from revengai.models.workflow_progress import WorkflowProgress from reai_toolkit.app.core.shared_schema import GenericApiReturn @@ -69,6 +71,22 @@ def _comments(items=None, status=TaskStatus.COMPLETED.value) -> CommentsData: return CommentsData.model_construct(inline_comments=items or [], task_status=status) +def _tokd(status=TaskStatus.COMPLETED.value) -> TokenisedData: + return TokenisedData.model_construct( + status=status, + tokenised_decompilation="int @@F@@(void) {}", + predicted_function_name="f", + function_mapping=MagicMock(), + ) + + +def _wait_mock(mock: MagicMock, timeout: float = 5.0) -> None: + deadline = time.monotonic() + timeout + while not mock.called and time.monotonic() < deadline: + time.sleep(0.005) + assert mock.called, "callback was not invoked" + + def _wait(service, timeout: float = 5.0) -> None: deadline = time.monotonic() + timeout while service.is_worker_running() and time.monotonic() < deadline: @@ -346,3 +364,144 @@ def stall_status(*a, **kw): release.set() time.sleep(0.2) on_decomp.assert_not_called() + + +def test_tokenised_phase_caches_and_dispatches_on_completed(service, sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + sdk.get_ai_decompilation_tokenised.return_value = _tokd() + + on_tokenised = MagicMock() + service.start_ai_decomp_task( + ea=4096, + on_decomp=MagicMock(), + on_summary=MagicMock(), + on_comments=MagicMock(), + on_tokenised=on_tokenised, + ) + _wait(service) + + on_tokenised.assert_called_once() + assert on_tokenised.call_args[0][0].success is True + assert service._tokenised_cache[42] is sdk.get_ai_decompilation_tokenised.return_value + + +def test_tokenised_phase_skipped_when_no_callback(service, sdk): + sdk.get_ai_decompilation.return_value = _dd(code="ok") + _run(service) + sdk.get_ai_decompilation_tokenised.assert_not_called() + + +def test_apply_overrides_sends_body_refetches_and_caches(service, sdk): + sdk.upsert_ai_decompilation_overrides.return_value = MagicMock() + sdk.get_ai_decompilation.return_value = _dd(code="renamed") + sdk.get_ai_decompilation_tokenised.return_value = _tokd() + + on_decomp, on_tokenised = MagicMock(), MagicMock() + service.apply_overrides( + ea=4096, + overrides={"@@V_v5@@": "buf"}, + on_decomp=on_decomp, + on_tokenised=on_tokenised, + ) + _wait_mock(on_decomp) + _wait_mock(on_tokenised) + + _, kwargs = sdk.upsert_ai_decompilation_overrides.call_args + assert kwargs["function_id"] == 42 + assert kwargs["upsert_overrides_input_body"].overrides == {"@@V_v5@@": "buf"} + + payload = on_decomp.call_args[0][0] + assert payload.success is True + assert payload.data.decompilation == "renamed" + assert service._decomp_cache[42].decompilation == "renamed" + assert service._tokenised_cache[42] is sdk.get_ai_decompilation_tokenised.return_value + + +def test_apply_overrides_api_error_surfaces(service, sdk): + sdk.upsert_ai_decompilation_overrides.side_effect = ApiException(status=500, reason="x") + + on_decomp, on_tokenised = MagicMock(), MagicMock() + service.apply_overrides( + ea=4096, overrides={"a": "b"}, on_decomp=on_decomp, on_tokenised=on_tokenised + ) + _wait_mock(on_decomp) + + assert on_decomp.call_args[0][0].success is False + sdk.get_ai_decompilation.assert_not_called() + + +def test_set_comment_calls_patch_and_updates_cache(service, sdk): + updated = _comments(items=[InlineComment.model_construct(comment="hi", line=3)]) + sdk.patch_ai_decompilation_inline_comment.return_value = updated + + on_result = MagicMock() + service.set_comment(ea=4096, line=3, comment="hi", on_result=on_result) + _wait_mock(on_result) + + _, kwargs = sdk.patch_ai_decompilation_inline_comment.call_args + assert kwargs["function_id"] == 42 + assert kwargs["patch_comment_body"].comment == "hi" + assert kwargs["patch_comment_body"].line == 3 + assert on_result.call_args[0][0].success is True + assert service._comments_cache[42] is updated + + +def test_remove_comment_calls_delete_and_updates_cache(service, sdk): + updated = _comments(items=[]) + sdk.delete_ai_decompilation_inline_comment.return_value = updated + + on_result = MagicMock() + service.remove_comment(ea=4096, line=3, on_result=on_result) + _wait_mock(on_result) + + _, kwargs = sdk.delete_ai_decompilation_inline_comment.call_args + assert kwargs["function_id"] == 42 + assert kwargs["line"] == 3 + assert on_result.call_args[0][0].success is True + assert service._comments_cache[42] is updated + + +def test_comment_mutation_api_error_surfaces(service, sdk): + sdk.patch_ai_decompilation_inline_comment.side_effect = ApiException(status=403) + + on_result = MagicMock() + service.set_comment(ea=4096, line=1, comment="x", on_result=on_result) + _wait_mock(on_result) + + assert on_result.call_args[0][0].success is False + + +def test_mutation_unknown_function_id_reports_failure(service, sdk, netstore): + netstore.get_function_mapping.return_value.inverse_function_map = {} + on_result = MagicMock() + service.set_comment(ea=4096, line=1, comment="x", on_result=on_result) + on_result.assert_called_once() + assert on_result.call_args[0][0].success is False + sdk.patch_ai_decompilation_inline_comment.assert_not_called() + + +def test_invalidate_clears_all_caches_and_inflight(service): + service._decomp_cache[42] = object() + service._summary_cache[42] = object() + service._comments_cache[42] = object() + service._tokenised_cache[42] = object() + evt = threading.Event() + service._inflight[42] = evt + + service.invalidate(42) + + assert 42 not in service._decomp_cache + assert 42 not in service._summary_cache + assert 42 not in service._comments_cache + assert 42 not in service._tokenised_cache + assert 42 not in service._inflight + assert evt.is_set() + + +def test_function_id_for_and_invalidate_ea(service): + assert service.function_id_for(4096) == 42 + assert service.function_id_for(9999) is None + + service._decomp_cache[42] = object() + service.invalidate_ea(4096) + assert 42 not in service._decomp_cache diff --git a/tests/unit/auto_unstrip_status/test_sdk_schemas.py b/tests/unit/auto_unstrip_status/test_sdk_schemas.py new file mode 100644 index 00000000..ca4a45da --- /dev/null +++ b/tests/unit/auto_unstrip_status/test_sdk_schemas.py @@ -0,0 +1,5 @@ +from revengai import AutoUnstripStatusOutputBody + + +def test_auto_unstrip_status_output_body_fields(): + assert "status" in AutoUnstripStatusOutputBody.model_fields diff --git a/tests/unit/auto_unstrip_status/test_service.py b/tests/unit/auto_unstrip_status/test_service.py new file mode 100644 index 00000000..8dda53cf --- /dev/null +++ b/tests/unit/auto_unstrip_status/test_service.py @@ -0,0 +1,107 @@ +import threading +from unittest.mock import MagicMock + +import pytest +from revengai import ApiException +from revengai.models.auto_unstrip_status_output_body import ( + AutoUnstripStatusOutputBody, +) + +from reai_toolkit.app.services.auto_unstrip_status import ( + auto_unstrip_status as svc_mod, +) +from reai_toolkit.app.services.auto_unstrip_status.auto_unstrip_status import ( + AutoUnstripStatusService, +) + +ANALYSIS_ID = 4242 + + +@pytest.fixture(autouse=True) +def no_sleep(monkeypatch): + monkeypatch.setattr(svc_mod.time, "sleep", lambda *_: None) + + +@pytest.fixture +def netstore(): + return MagicMock() + + +@pytest.fixture +def service(netstore): + return AutoUnstripStatusService(netstore_service=netstore, sdk_config=MagicMock()) + + +@pytest.fixture +def sdk(mocker): + mocker.patch.object(AutoUnstripStatusService, "yield_api_client") + api_class = mocker.patch.object(svc_mod, "AnalysesCoreApi") + api_inst = MagicMock() + api_class.return_value = api_inst + return api_inst + + +def _status(value): + return AutoUnstripStatusOutputBody.model_construct(status=value) + + +def _poll(service, resync_if_already_complete=True): + cb = MagicMock() + service._thread_callback = cb + service._poll_auto_unstrip_status( + threading.Event(), ANALYSIS_ID, resync_if_already_complete + ) + return cb + + +def test_completed_requests_resync(service, sdk): + sdk.v3_get_analysis_auto_unstrip_status.return_value = _status("COMPLETED") + + cb = _poll(service, resync_if_already_complete=True) + + cb.assert_called_once() + result = cb.call_args[0][0] + assert result.success is True + assert result.data == ANALYSIS_ID + + +def test_already_complete_skips_resync_when_attaching(service, sdk): + sdk.v3_get_analysis_auto_unstrip_status.return_value = _status("COMPLETED") + + cb = _poll(service, resync_if_already_complete=False) + + cb.assert_not_called() + + +def test_failed_status_reports_failure(service, sdk): + sdk.v3_get_analysis_auto_unstrip_status.return_value = _status("FAILED") + + cb = _poll(service) + + result = cb.call_args[0][0] + assert result.success is False + + +def test_running_keeps_polling_until_completed(service, sdk): + sdk.v3_get_analysis_auto_unstrip_status.side_effect = [ + _status("RUNNING"), + _status("COMPLETED"), + ] + + cb = _poll(service, resync_if_already_complete=False) + + assert sdk.v3_get_analysis_auto_unstrip_status.call_count == 2 + result = cb.call_args[0][0] + assert result.success is True + assert result.data == ANALYSIS_ID + + +def test_status_request_failure_surfaces_error(service, sdk): + sdk.v3_get_analysis_auto_unstrip_status.side_effect = ApiException( + status=500, reason="boom" + ) + + cb = _poll(service) + + result = cb.call_args[0][0] + assert result.success is False diff --git a/tests/unit/core/test_logging_setup.py b/tests/unit/core/test_logging_setup.py new file mode 100644 index 00000000..d4ed0474 --- /dev/null +++ b/tests/unit/core/test_logging_setup.py @@ -0,0 +1,45 @@ +from unittest.mock import MagicMock + +import pytest + +from reai_toolkit.app.core import logging_setup + + +@pytest.fixture +def fake_logger(monkeypatch): + logger = MagicMock() + monkeypatch.setattr(logging_setup, "logger", logger) + return logger + + +def _added_level(logger) -> str: + return logger.add.call_args.kwargs["level"] + + +def test_default_is_info(monkeypatch, fake_logger): + monkeypatch.delenv("REAI_DEBUG", raising=False) + monkeypatch.delenv("REAI_LOG_LEVEL", raising=False) + logging_setup.configure_logging() + fake_logger.remove.assert_called_once() + assert _added_level(fake_logger) == "INFO" + + +def test_reai_debug_enables_debug(monkeypatch, fake_logger): + monkeypatch.delenv("REAI_LOG_LEVEL", raising=False) + monkeypatch.setenv("REAI_DEBUG", "1") + logging_setup.configure_logging() + assert _added_level(fake_logger) == "DEBUG" + + +def test_reai_debug_zero_stays_info(monkeypatch, fake_logger): + monkeypatch.delenv("REAI_LOG_LEVEL", raising=False) + monkeypatch.setenv("REAI_DEBUG", "0") + logging_setup.configure_logging() + assert _added_level(fake_logger) == "INFO" + + +def test_explicit_level_wins_and_is_uppercased(monkeypatch, fake_logger): + monkeypatch.setenv("REAI_DEBUG", "1") + monkeypatch.setenv("REAI_LOG_LEVEL", "warning") + logging_setup.configure_logging() + assert _added_level(fake_logger) == "WARNING"