New feature: automatic variable renaming.

The API key can now be obtain from an envvar too.

README.md

@@ -1,6 +1,12 @@
+# 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.
+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)
 
-# Setup
+## Setup
 
 Simply drop this script into your IDA plugins folder (`$IDAUSR/plugins`).
 
@@ -16,18 +22,30 @@ 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.
 
-# Usage
+## Usage
 
-Once the plugin is installed properly, you should be able to invoke it from the context menu of IDA's pseudo code windows, as shown in the screenshot below:
+Once the plugin is installed properly, you should be able to invoke it from the context menu of IDA's pseudocode window,
+as shown in the screenshot below:
 
 ![](https://github.com/JusticeRage/Gepetto/blob/main/readme/usage.png?raw=true)
 
-# Limitations
+You can also use the following hotkeys:
+
+- Ask ChatGPT 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.
+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!
 
-# Acknowledgements
+## Acknowledgements
 
 - [OpenAI](https://openai.com), for making this incredible chatbot, obviously
+- [Hex Rays](https://hex-rays.com/), the makers of IDA for their lightning fast support
 - [Kaspersky](https://kaspersky.com), for funding all my research

gepetto.py

@@ -1,12 +1,16 @@
 import functools
+import json
 import idaapi
 import ida_hexrays
 import ida_kernwin
 import idc
 import openai
+import os
+import re
 import textwrap
 import threading
 
+# Set your API key here, or put in in the OPENAI_API_KEY environment variable.
 openai.api_key = ""
 
 # =============================================================================
@@ -15,29 +19,40 @@
 
 class GepettoPlugin(idaapi.plugin_t):
     flags = 0
-    action_name = "gepetto:explain_function"
-    menu_path = "Edit/Gepetto/Explain function"
+    explain_action_name = "gepetto:explain_function"
+    explain_menu_path = "Edit/Gepetto/Explain function"
+    rename_action_name = "gepetto:rename_function"
+    rename_menu_path = "Edit/Gepetto/Rename variables"
     wanted_name = 'Gepetto'
     wanted_hotkey = ''
+    comment = "Uses ChatGPT to enrich the decompiler's output"
+    help = "See usage instructions on GitHub"
     menu = None
 
     def init(self):
         # Check for whether the decompiler is available
         if not ida_hexrays.init_hexrays_plugin():
             return idaapi.PLUGIN_SKIP
 
-        action_desc = idaapi.action_desc_t(
-            self.action_name,               # The action name. This acts like an ID and must be unique
-            'Explain function',             # The action text.
-            GepettoHandler(),               # The action handler.
-            "Ctrl+Alt+G",                   # Optional: the action shortcut
-            'Use ChatGPT to explain the currently selected function',  # Optional: the action tooltip (available in menus/toolbar)
-            199)                            # Optional: the action icon (shows when in menus/toolbars)
-        idaapi.register_action(action_desc)
-        idaapi.attach_action_to_menu(
-            self.menu_path,                 # The relative path of where to add the action
-            self.action_name,               # The action ID (see above)
-            idaapi.SETMENU_APP)             # We want to append the action after the 'Manual instruction...'
+        # Function explaining action
+        explain_action = idaapi.action_desc_t(self.explain_action_name,
+                                              'Explain function',
+                                              ExplainHandler(),
+                                              "Ctrl+Alt+G",
+                                              'Use ChatGPT 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)
+
+        # Variable renaming action
+        rename_action = idaapi.action_desc_t(self.rename_action_name,
+                                             'Rename variables',
+                                             RenameHandler(),
+                                             "Ctrl+Alt+R",
+                                             "Use ChatGPT 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)
 
         # Register context menu actions
         self.menu = ContextMenuHooks()
@@ -49,13 +64,23 @@ def run(self, arg):
         pass
 
     def term(self):
-        idaapi.detach_action_from_menu(self.menu_path, self.action_name)
+        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)
         if self.menu:
             self.menu.unhook()
         return
 
 # -----------------------------------------------------------------------------
 
+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:
+            idaapi.attach_action_to_popup(form, popup, GepettoPlugin.explain_action_name, "Gepetto/")
+            idaapi.attach_action_to_popup(form, popup, GepettoPlugin.rename_action_name, "Gepetto/")
+
+# -----------------------------------------------------------------------------
+
 def comment_callback(address, view, response):
     """
     Callback that sets a comment at the given address.
@@ -75,15 +100,20 @@ def comment_callback(address, view, response):
 
 # -----------------------------------------------------------------------------
 
-class GepettoHandler(idaapi.action_handler_t):
+class ExplainHandler(idaapi.action_handler_t):
+    """
+    This handler is tasked with querying ChatGPT for an explanation of the
+    given function. Once the reply is received, it is added as a function
+    comment.
+    """
     def __init__(self):
         idaapi.action_handler_t.__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?\n" + str(decompiler_output) +
-                            "\nSuggest better names for the function and its arguments.",
+        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))
         return 1
 
@@ -93,11 +123,54 @@ def update(self, ctx):
 
 # -----------------------------------------------------------------------------
 
-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:
-            idaapi.attach_action_to_popup(form, popup, GepettoPlugin.action_name, "Gepetto/")
+def rename_callback(address, view, response):
+    """
+    Callback that extracts a JSON array of old names and new names from the
+    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
+    """
+    j = re.search(r"\{[^}]*?\}", response)
+    if not j:
+        print(f"Error: couldn't extract a response from ChatGPT'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
+    for n in names:
+        if ida_hexrays.rename_lvar(function_addr, n, names[n]):
+            counter += 1
+    # 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.")
+
+# -----------------------------------------------------------------------------
+
+class RenameHandler(idaapi.action_handler_t):
+    """
+    This handler requests new variable names from ChatGPT and updates the
+    decompiler's output.
+    """
+    def __init__(self):
+        idaapi.action_handler_t.__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) +
+                            "\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))
+        return 1
+
+    # This action is always available.
+    def update(self, ctx):
+        return idaapi.AST_ENABLE_ALWAYS
 
 # =============================================================================
 # ChatGPT interaction
@@ -124,6 +197,7 @@ def query_chatgpt(query, cb):
     except openai.OpenAIError as e:
         raise print(f"ChatGPT could not complete the request: {str(e)}")
 
+# -----------------------------------------------------------------------------
 
 def query_chatgpt_async(query, cb):
     """
@@ -141,7 +215,9 @@ def query_chatgpt_async(query, cb):
 
 def PLUGIN_ENTRY():
     if not openai.api_key:
-        print("Please edit this script to insert your OpenAI API key!")
-        raise ValueError("No valid OpenAI API key found")
+        openai.api_key = os.getenv("OPENAI_API_KEY")
+        if not openai.api_key:
+            print("Please edit this script to insert your OpenAI API key!")
+            raise ValueError("No valid OpenAI API key found")
 
     return GepettoPlugin()