When renaming variables, also update mentions of them in the function…

… comment.

Changed all references of ChatGPT to davinci-003 (closes JusticeRage#4)
Minor bugfixes (possibly closes JusticeRage#5)

README.md

@@ -1,7 +1,7 @@
 # Gepetto
 
-Gepetto is a Python script which uses OpenAI's ChatGPT to provide meaning to functions decompiled by IDA Pro.
-At the moment, it can ask ChatGPT to explain what a function does, and to automatically rename its variables.
+Gepetto is a Python script which uses OpenAI's davinci-003 model to provide meaning to functions decompiled by IDA Pro.
+At the moment, it can ask davinci-003 to explain what a function does, and to automatically rename its variables.
 Here is a simple example of what results it can provide in mere seconds:
 
 ![](https://github.com/JusticeRage/Gepetto/blob/main/readme/comparison.png?raw=true)
@@ -20,7 +20,7 @@ Finally, with the corresponding interpreter, simply run:
 ```
 
 ⚠️ You will also need to edit the script and add your own API key, which can be found on [this page](https://beta.openai.com/account/api-keys).
-Please note that ChatGPT queries are not free (although not very expensive) and you will need to setup a payment method.
+Please note that davinci-003 queries are not free (although not very expensive) and you will need to setup a payment method.
 
 ## Usage
 
@@ -31,18 +31,18 @@ as shown in the screenshot below:
 
 You can also use the following hotkeys:
 
-- Ask ChatGPT to explain the function: `Ctrl` + `Alt` + `H`
+- Ask davinci-003 to explain the function: `Ctrl` + `Alt` + `H`
 - Request better names for the function's variables: `Ctrl` + `Alt` + `R`
 
 Initial testing shows that asking for better names works better if you ask for an explanation of the function first – I
-assume because ChatGPT then uses its own comment to make more accurate suggestions.
+assume because davinci-003 then uses its own comment to make more accurate suggestions.
 There is an element of randomness to the AI's replies. If for some reason the initial response you get doesn't suit you,
 you can always run the command again.
 
 ## Limitations
 
 - The plugin requires access to the HexRays decompiler to function.
-- ChatGPT is a general-purpose chatbot and may very well get things wrong! Always be critical of results returned!
+- davinci-003 is a general-purpose language model and may very well get things wrong! Always be critical of results returned!
 
 ## Acknowledgements
 

gepetto.py

@@ -25,12 +25,12 @@ class GepettoPlugin(idaapi.plugin_t):
     rename_menu_path = "Edit/Gepetto/Rename variables"
     wanted_name = 'Gepetto'
     wanted_hotkey = ''
-    comment = "Uses ChatGPT to enrich the decompiler's output"
+    comment = "Uses davinci-003 to enrich the decompiler's output"
     help = "See usage instructions on GitHub"
     menu = None
 
     def init(self):
-        # Check for whether the decompiler is available
+        # Check whether the decompiler is available
         if not ida_hexrays.init_hexrays_plugin():
             return idaapi.PLUGIN_SKIP
 
@@ -39,7 +39,7 @@ def init(self):
                                               'Explain function',
                                               ExplainHandler(),
                                               "Ctrl+Alt+G",
-                                              'Use ChatGPT to explain the currently selected function',
+                                              'Use davinci-003 to explain the currently selected function',
                                               199)
         idaapi.register_action(explain_action)
         idaapi.attach_action_to_menu(self.explain_menu_path, self.explain_action_name, idaapi.SETMENU_APP)
@@ -49,7 +49,7 @@ def init(self):
                                              'Rename variables',
                                              RenameHandler(),
                                              "Ctrl+Alt+R",
-                                             "Use ChatGPT to rename this function's variables",
+                                             "Use davinci-003 to rename this function's variables",
                                              199)
         idaapi.register_action(rename_action)
         idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)
@@ -65,7 +65,7 @@ def run(self, arg):
 
     def term(self):
         idaapi.detach_action_from_menu(self.explain_menu_path, self.explain_action_name)
-        idaapi.detach_action_from_menu(self.explain_menu_path, self.rename_action_name)
+        idaapi.detach_action_from_menu(self.rename_menu_path, self.rename_action_name)
         if self.menu:
             self.menu.unhook()
         return
@@ -75,7 +75,7 @@ def term(self):
 class ContextMenuHooks(idaapi.UI_Hooks):
     def finish_populating_widget_popup(self, form, popup):
         # Add actions to the context menu of the Pseudocode view
-        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE or idaapi.get_widget_type(form) == idaapi.BWN_DISASM:
+        if idaapi.get_widget_type(form) == idaapi.BWN_PSEUDOCODE:
             idaapi.attach_action_to_popup(form, popup, GepettoPlugin.explain_action_name, "Gepetto/")
             idaapi.attach_action_to_popup(form, popup, GepettoPlugin.rename_action_name, "Gepetto/")
 
@@ -94,15 +94,16 @@ def comment_callback(address, view, response):
     # Add the response as a comment in IDA.
     idc.set_func_cmt(address, response, 0)
     # Refresh the window so the comment is displayed properly
-    view.refresh_view(False)
-    print("ChatGPT query finished!")
+    if view:
+        view.refresh_view(False)
+    print("davinci-003 query finished!")
 
 
 # -----------------------------------------------------------------------------
 
 class ExplainHandler(idaapi.action_handler_t):
     """
-    This handler is tasked with querying ChatGPT for an explanation of the
+    This handler is tasked with querying davinci-003 for an explanation of the
     given function. Once the reply is received, it is added as a function
     comment.
     """
@@ -112,9 +113,9 @@ def __init__(self):
     def activate(self, ctx):
         decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
         v = ida_hexrays.get_widget_vdui(ctx.widget)
-        query_chatgpt_async("Can you explain what the following C function does and suggest a better name for it?\n"
-                            + str(decompiler_output),
-                            functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
+        query_model_async("Can you explain what the following C function does and suggest a better name for it?\n"
+                          + str(decompiler_output),
+                          functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
         return 1
 
     # This action is always available.
@@ -129,30 +130,39 @@ def rename_callback(address, view, response):
     response and sets them in the pseudocode.
     :param address: The address of the function to work on
     :param view: A handle to the decompiler window
-    :param response: The response from ChatGPT
+    :param response: The response from davinci-003
     """
     j = re.search(r"\{[^}]*?\}", response)
     if not j:
-        print(f"Error: couldn't extract a response from ChatGPT's output:\n{response}")
+        print(f"Error: couldn't extract a response from davinci-003's output:\n{response}")
         return
     names = json.loads(j.group(0))
 
     # The rename function needs the start address of the function
     function_addr = idaapi.get_func(address).start_ea
 
-    counter = 0
+    replaced = []
     for n in names:
         if ida_hexrays.rename_lvar(function_addr, n, names[n]):
-            counter += 1
+            replaced.append(n)
+
+    # Update possible names left in the function comment
+    comment = idc.get_func_cmt(address, 0)
+    if comment and len(replaced) > 0:
+        for n in replaced:
+            comment = re.sub(r'\b%s\b' % n, names[n], comment)
+        idc.set_func_cmt(address, comment, 0)
+
     # Refresh the window to show the new names
-    view.refresh_view(True)
-    print(f"ChatGPT query finished! {counter} variable(s) renamed.\nPress F5 if the new names don't appear.")
+    if view:
+        view.refresh_view(True)
+    print(f"davinci-003 query finished! {len(replaced)} variable(s) renamed.")
 
 # -----------------------------------------------------------------------------
 
 class RenameHandler(idaapi.action_handler_t):
     """
-    This handler requests new variable names from ChatGPT and updates the
+    This handler requests new variable names from davinci-003 and updates the
     decompiler's output.
     """
     def __init__(self):
@@ -161,26 +171,26 @@ def __init__(self):
     def activate(self, ctx):
         decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
         v = ida_hexrays.get_widget_vdui(ctx.widget)
-        query_chatgpt_async("Analyze the following C function:\n" + str(decompiler_output) +
+        query_model_async("Analyze the following C function:\n" + str(decompiler_output) +
                             "\nSuggest better variable names, reply with a JSON array where keys are the original names"
                             "and values are the proposed names. Do not explain anything, only print the JSON "
                             "dictionary.",
-                            functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=v))
+                          functools.partial(rename_callback, address=idaapi.get_screen_ea(), view=v))
         return 1
 
     # This action is always available.
     def update(self, ctx):
         return idaapi.AST_ENABLE_ALWAYS
 
 # =============================================================================
-# ChatGPT interaction
+# davinci-003 interaction
 # =============================================================================
 
-def query_chatgpt(query, cb):
+def query_model(query, cb):
     """
-    Function which sends a query to ChatGPT and calls a callback when the response is available.
+    Function which sends a query to davinci-003 and calls a callback when the response is available.
     Blocks until the response is received
-    :param query: The request to send to ChatGPT
+    :param query: The request to send to davinci-003
     :param cb: Tu function to which the response will be passed to.
     """
     try:
@@ -191,22 +201,25 @@ def query_chatgpt(query, cb):
             max_tokens=2500,
             top_p=1,
             frequency_penalty=1,
-            presence_penalty=1
+            presence_penalty=1,
+            timeout=60  # Wait 60 seconds maximum
         )
         ida_kernwin.execute_sync(functools.partial(cb, response=response.choices[0].text), ida_kernwin.MFF_WRITE)
     except openai.OpenAIError as e:
-        raise print(f"ChatGPT could not complete the request: {str(e)}")
+        print(f"davinci-003 could not complete the request: {str(e)}")
+    except Exception as e:
+        print(f"General exception encountered while running the query: {str(e)}")
 
 # -----------------------------------------------------------------------------
 
-def query_chatgpt_async(query, cb):
+def query_model_async(query, cb):
     """
-    Function which sends a query to ChatGPT and calls a callback when the response is available.
-    :param query: The request to send to ChatGPT
+    Function which sends a query to davinci-003 and calls a callback when the response is available.
+    :param query: The request to send to davinci-003
     :param cb: Tu function to which the response will be passed to.
     """
-    print("Request to ChatGPT sent...")
-    t = threading.Thread(target=query_chatgpt, args=[query, cb])
+    print("Request to davinci-003 sent...")
+    t = threading.Thread(target=query_model, args=[query, cb])
     t.start()
 
 # =============================================================================