add watching args onLeave feature

Author: hackcatml

code.py

@@ -203,8 +203,8 @@ def set_watch(self, addr, is_reg_watch):
     def detach_all(self):
         self.script.exports.detachall()
 
-    def set_read_args_options(self, addr, index, option):
-        self.script.exports.setreadargsoptions(addr, index, option)
+    def set_read_args_options(self, addr, index, option, on_leave):
+        self.script.exports.setreadargsoptions(addr, index, option, on_leave)
 
     def set_read_retval_options(self, addr, option):
         self.script.exports.setreadretvalsoptions(addr, option)

hexviewer.py

@@ -5,7 +5,7 @@
 from PyQt6.QtCore import Qt, pyqtSlot
 from PyQt6.QtGui import QTextCursor, QAction, QCursor
 from PyQt6.QtWidgets import QTextEdit, QApplication, QWidget, QVBoxLayout, QSlider, QLabel, QHBoxLayout, QListWidget, \
-    QPushButton, QMenu
+    QPushButton, QMenu, QCheckBox, QWidgetAction
 
 import globvar
 
@@ -336,9 +336,35 @@ def __init__(self, text):
 
 # Custom TextEdit class for NewWatchWidget
 class CustomTextEdit(QTextEdit):
+    def __init__(self, parent=None):
+        super(CustomTextEdit, self).__init__(parent)
+        self.key = None
+        self.args_index = None
+        self.addr = None
+        self.checkedStates = {}
+
+    def closeEvent(self, e: QtGui.QCloseEvent) -> None:
+        self.checkedStates.clear()
+
+    def get_args_index_and_addr_from_selected_text(self):
+        # Get the currently selected text
+        tc = self.textCursor()
+        selected_text = self.textCursor().selectedText()
+        if ":" in selected_text:
+            self.args_index = selected_text[4:selected_text.index(":")]
+        else:
+            self.args_index = selected_text[4:]
+
+        while True:
+            tc.movePosition(QTextCursor.MoveOperation.Up, QTextCursor.MoveMode.MoveAnchor)
+            if re.search(r"] 0x[a-f0-9]+", tc.block().text()):
+                self.addr = tc.block().text()[4:].strip()
+                break
+
     def contextMenuEvent(self, e: QtGui.QContextMenuEvent) -> None:
         menu = super(CustomTextEdit, self).createStandardContextMenu()  # Get the default context menu
         select_all_action = None
+        self.get_args_index_and_addr_from_selected_text()
 
         for action in menu.actions():  # loop over the existing actions
             if action.text() == "Select All":
@@ -349,6 +375,21 @@ def contextMenuEvent(self, e: QtGui.QContextMenuEvent) -> None:
             args_regx = re.compile(r'(\bargs\d+\b|\breturn\b)')
             match = args_regx.match(self.textCursor().selectedText())
 
+            on_leave_check = QCheckBox("OnLeave", self)
+            if match is not None:
+                self.key = (self.addr, self.args_index)  # Use a tuple as the key
+                # If the key is not in the checkedStates dictionary, add it with a default state of unchecked
+                if self.key not in self.checkedStates:
+                    self.checkedStates[self.key] = Qt.CheckState.Unchecked
+                # Set the checked state based on the checkedStates dictionary
+                on_leave_check.setCheckState(Qt.CheckState(self.checkedStates[self.key]))
+                # Connect the checkbox's stateChanged signal to a function
+                on_leave_check.stateChanged.connect(lambda state: self.on_leave_check_state_changed(state, self.key))
+
+                check_action = QWidgetAction(self)
+                check_action.setDefaultWidget(on_leave_check)
+                menu.insertAction(select_all_action, check_action)
+
             read_pointer_action = QAction("readPointer", self)
             if match is not None:
                 read_pointer_action.setEnabled(True)
@@ -394,6 +435,11 @@ def contextMenuEvent(self, e: QtGui.QContextMenuEvent) -> None:
             # Show the context menu.
             menu.exec(e.globalPos())
 
+    # handle the checkbox state change
+    def on_leave_check_state_changed(self, state, key):
+        # Update the checked state in the checkedStates dictionary
+        self.checkedStates[key] = state
+
     def read_pointer(self):
         self.read_args_with_options("readPointer")
 
@@ -416,18 +462,12 @@ def reset(self):
         self.read_args_with_options("")
 
     def read_args_with_options(self, option):
-        tc = self.textCursor()
-        args_index = self.textCursor().selectedText()[4:]
-        while True:
-            tc.movePosition(QTextCursor.MoveOperation.Up, QTextCursor.MoveMode.MoveAnchor)
-            if re.search(r"] 0x[a-f0-9]+", tc.block().text()):
-                addr = tc.block().text()[4:].strip()
-                break
         try:
             if self.textCursor().selectedText() == "return":
-                globvar.fridaInstrument.set_read_retval_options(addr, option)
+                globvar.fridaInstrument.set_read_retval_options(self.addr, option)
             else:
-                globvar.fridaInstrument.set_read_args_options(addr, args_index, option)
+                globvar.fridaInstrument.set_read_args_options(self.addr, self.args_index, option,
+                                                              self.checkedStates[self.key])
         except Exception as e:
             print(f"{inspect.currentframe().f_code.co_name}: {e}")
             return
@@ -470,6 +510,7 @@ def __init__(self):
 
     def closeEvent(self, e: QtGui.QCloseEvent) -> None:
         self.text_edit.clear()
+        self.text_edit.closeEvent(e)
         self.watch_list.clear()
         try:
             globvar.fridaInstrument.detach_all()

scripts/default.js

@@ -37,12 +37,15 @@ function readargs(args, index, addr) {
     // if the argument's index is in the read_args_options for the current address
     if (read_args_options[addr] && read_args_options[addr][index]) {
         // get the option from the read_args_options
-        var option = read_args_options[addr][index];
+        var option = read_args_options[addr][index]["readOption"];
         // read the argument based on the option
         if(option === "readByteArray") {
             let data = new Uint8Array(args[option](32))
             return Object.keys(data).map(key => data[key].toString(16).padStart(2, '0')).join(' ');
         }
+        if(option === '') {
+            return args
+        }
         return args[option]();
     } else {
         // just log the argument if it's not in the read_args_options
@@ -278,13 +281,15 @@ rpc.exports = {
         count = 0
         modules.length = 0
     },
-    setreadargsoptions: (addr, index, option) => {
+    setreadargsoptions: (addr, index, option, onleave) => {
         // if the address is not in the read_args_options yet, add it
         if (!read_args_options[addr]) {
             read_args_options[addr] = {};
         }
-        // add/update the index and option for the address
-        read_args_options[addr][index] = option;
+        // add/update the index and read options for the address
+        read_args_options[addr][index] = {};
+        read_args_options[addr][index]["readOption"] = option
+        read_args_options[addr][index]["onLeave"] = onleave !== 0;
     },
     setreadretvalsoptions: (addr, option) => {
         // if the address is not in the read_args_options yet, add it
@@ -300,6 +305,11 @@ rpc.exports = {
     setwatch: (addr, is_reg_watch) => {
         Interceptor.attach(ptr(addr), {
             onEnter: function (args) {
+                this.argv = []
+                for (let index = 0; index < num_args_to_watch; index++) {
+                    this.argv[index] = args[index]
+                }
+
                 if(is_reg_watch) {
                     let count = 0
                     let log = `[+] ${ptr(addr)}:\n`
@@ -317,10 +327,17 @@ rpc.exports = {
                 } else {
                     let log = `[+] ${ptr(addr)}\n`
                     for (let index = 0; index < num_args_to_watch - 1; index++) {
+                        if(read_args_options[ptr(addr)] && read_args_options[ptr(addr)][index] && read_args_options[ptr(addr)][index]["onLeave"]) {
+                            continue;
+                        }
                         log += `args${index}: ${readargs(args[index], index, ptr(addr))}, `
                     }
-                    log += `args${num_args_to_watch - 1}: ${readargs(args[num_args_to_watch - 1], num_args_to_watch - 1, ptr(addr))}`
-                    send({'watchArgs':log})
+                    if(read_args_options[ptr(addr)] && read_args_options[ptr(addr)][num_args_to_watch - 1] && read_args_options[ptr(addr)][num_args_to_watch - 1]["onLeave"]) {
+                        send({'watchArgs':log})
+                    } else {
+                        log += `args${num_args_to_watch - 1}: ${readargs(args[num_args_to_watch - 1], num_args_to_watch - 1, ptr(addr))}`
+                        send({'watchArgs':log})
+                    }
                 }
             },
             onLeave: function(retval) {
@@ -340,6 +357,15 @@ rpc.exports = {
                     send({'watchRegs':log})
                 } else {
                     let log = `return: ${readretval(retval, ptr(addr))}\n`
+                    for (let index = 0; index < num_args_to_watch; index++) {
+                        if(read_args_options[ptr(addr)] && read_args_options[ptr(addr)][index] && read_args_options[ptr(addr)][index]["onLeave"]){
+                            log += `args${index}: ${readargs(this.argv[index], index, ptr(addr))}, `
+                            this.onleave = true
+                        }
+                    }
+                    if (this.onleave) {
+                        log += '\n'
+                    }
                     log += `[-] ${ptr(addr)}`
                     send({'watchArgs':log})
                 }