Add feature: An option to rewrite code

README.md

@@ -48,6 +48,7 @@ You can also use the following hotkeys:
 
 - Ask the model to explain the function: `Ctrl` + `Alt` + `H`
 - Request better names for the function's variables: `Ctrl` + `Alt` + `R`
+- Request code rewrite for the function's variables: `Ctrl` + `Alt` + `W`
 
 Initial testing shows that asking for better names works better if you ask for an explanation of the function first – I
 assume because the model then uses its own comment to make more accurate suggestions.

gepetto/ida/handlers.py

@@ -57,8 +57,8 @@ def activate(self, ctx):
         decompiler_output = ida_hexrays.decompile(idaapi.get_screen_ea())
         v = ida_hexrays.get_widget_vdui(ctx.widget)
         gepetto.config.model.query_model_async(
-            _("Can you explain what the following C function does and suggest a better name for "
-              "it?\n{decompiler_output}").format(decompiler_output=str(decompiler_output)),
+            _("Can you explain what the following C function does? "
+              "\n{decompiler_output}").format(decompiler_output=str(decompiler_output)),
             functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
         return 1
 
@@ -134,6 +134,33 @@ def activate(self, ctx):
     # This action is always available.
     def update(self, ctx):
         return idaapi.AST_ENABLE_ALWAYS
+
+
+# -----------------------------------------------------------------------------
+
+class RewriteHandler(idaapi.action_handler_t):
+    """
+    This handler is tasked with querying the model for a rewrite of the
+    given code. 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)
+        gepetto.config.model.query_model_async(
+            _("Can you rewrite a better code from this following C function? "
+              "\n{decompiler_output}\n"
+              "Do not explain anything, only print your rewriten code. Optimizaton are appreciated.").format(decompiler_output=str(decompiler_output)),
+            functools.partial(comment_callback, address=idaapi.get_screen_ea(), view=v))
+        return 1
+
+    # This action is always available.
+    def update(self, ctx):
+        return idaapi.AST_ENABLE_ALWAYS
 
 
 # -----------------------------------------------------------------------------

gepetto/ida/ui.py

@@ -5,7 +5,7 @@
 import ida_hexrays
 
 import gepetto.config
-from gepetto.ida.handlers import ExplainHandler, RenameHandler, SwapModelHandler
+from gepetto.ida.handlers import ExplainHandler, RenameHandler, RewriteHandler, SwapModelHandler
 from gepetto.models.base import GPT4_MODEL_NAME, GPT3_MODEL_NAME
 
 _ = gepetto.config.translate.gettext
@@ -20,6 +20,8 @@ class GepettoPlugin(idaapi.plugin_t):
     explain_menu_path = "Edit/Gepetto/" + _("Explain function")
     rename_action_name = "gepetto:rename_function"
     rename_menu_path = "Edit/Gepetto/" + _("Rename variables")
+    rewrite_action_name = "gepetto:rewrite_function"
+    rewrite_menu_path = "Edit/Gepetto/" + _("Rewrite code")
 
     # Model selection menu
     select_gpt35_action_name = "gepetto:select_gpt35"
@@ -61,6 +63,17 @@ def init(self):
                                              201)
         idaapi.register_action(rename_action)
         idaapi.attach_action_to_menu(self.rename_menu_path, self.rename_action_name, idaapi.SETMENU_APP)
+
+        # Variable rewrite action
+        rewrite_action = idaapi.action_desc_t(self.rewrite_action_name,
+                                             _('Rewrite code'),
+                                             RewriteHandler(),
+                                             "Ctrl+Alt+W",
+                                             _("Use {model} to rewrite this function's code").format(
+                                                 model=str(gepetto.config.model)),
+                                             201)
+        idaapi.register_action(rewrite_action)
+        idaapi.attach_action_to_menu(self.rewrite_menu_path, self.rewrite_action_name, idaapi.SETMENU_APP)
 
         self.generate_plugin_select_menu()
 
@@ -118,6 +131,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.rename_menu_path, self.rename_action_name)
+        idaapi.detach_action_from_menu(self.rename_menu_path, self.rewrite_action_name)
         idaapi.detach_action_from_menu(self.select_gpt35_menu_path, self.select_gpt35_action_name)
         idaapi.detach_action_from_menu(self.select_gpt4_menu_path, self.select_gpt4_action_name)
         if self.menu:
@@ -132,3 +146,4 @@ def finish_populating_widget_popup(self, form, popup):
         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/")
+            idaapi.attach_action_to_popup(form, popup, GepettoPlugin.rewrite_action_name, "Gepetto/")