Implement renaming individual variables.

- Works only with HLIL Variable Initializations. Assignments would
  also be appropriate, but I found that the model typically only
  returned "result" as the string. This was not useful.

Author: WhatTheFuzz

__init__.py

@@ -1,7 +1,7 @@
 from binaryninja import PluginCommand
 from . src.settings import OpenAISettings
-from . src.entry import check_function
-from . src.entry import rename_expression
+from . src.entry import check_function, rename_variable, \
+                        rename_all_variables_in_function
 
 # Register the settings group in Binary Ninja to store the API key and model.
 OpenAISettings()
@@ -20,10 +20,13 @@
                             "OPENAI_API_KEY or modify the path in entry.py.",
                             check_function)
 
-PluginCommand.register_for_high_level_il_instruction("OpenAI\Rename Variable",
+PluginCommand.register_for_high_level_il_instruction("OpenAI\Rename Variable (HLIL)",
                             "If the current expression is a HLIL Initialization " \
                             "(HighLevelILVarInit), then query OpenAI to rename the " \
                             "variable to what it believes is correct. If the expression" \
                             "is not an HighLevelILVarInit, then do nothing. Requires " \
                             "an internet connection and an API key. ",
-                            rename_expression)
+                            rename_variable)
+
+PluginCommand.register_for_high_level_il_function("OpenAI\Rename All Variables (HLIL)",
+"", rename_all_variables_in_function)

src/agent.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 from pathlib import Path
 
 import openai
@@ -9,7 +9,8 @@
 from binaryninja.function import Function
 from binaryninja.lowlevelil import LowLevelILFunction
 from binaryninja.mediumlevelil import MediumLevelILFunction
-from binaryninja.highlevelil import HighLevelILFunction, HighLevelILInstruction
+from binaryninja.highlevelil import HighLevelILFunction, HighLevelILInstruction, \
+                                    HighLevelILVarInit
 from binaryninja.settings import Settings
 from binaryninja import log, BinaryView
 
@@ -24,10 +25,10 @@ class Agent:
     It is in IL_FORM. What does this function do?
     '''
 
-    rename_variable_question: str = '''
-    In one word, what should the variable name be for the variable that is
-    assigned to the result of the C expression:\n
-    '''
+    rename_variable_question: str = "In one word, what should the variable " \
+        "be for the variable that is assigned to the result of the C " \
+        "expression:\n"
+
 
     # A mapping of IL forms to their names.
     il_name: dict[type, str] = {
@@ -48,6 +49,8 @@ def __init__(self,
         # Set instance attributes.
         self.bv = bv
         self.model = self.get_model()
+        # Used for the callback function.
+        self.instruction = None
 
     def read_api_key(self, filename: Optional[Path]=None) -> str:
         '''Checks for the API key in three locations.
@@ -157,16 +160,46 @@ def generate_query(self, function: Union[Function,
         prompt += '\n'.join(self.instruction_list(function))
         return prompt
 
-    def generate_rename_expression_query(
-                                    instruction: HighLevelILInstruction) -> str:
+    def generate_rename_variable_query(self,
+                                    instruction: HighLevelILInstruction) -> Optional[str]:
         '''Generates a query string given a BNIL instruction. Returns the query
         as a string.
         '''
-        pass
+        if not isinstance(instruction, HighLevelILVarInit):
+            raise TypeError(f'Expected a BNIL instruction of type '
+                            f'HighLevelILVarInit got {type(instruction)}.')
+        # Assign the instruction to the Agent instance. This is used for the
+        # callback function so we don't need to pass in the instruction to the
+        # Query instance. This is kind of janky and should be examined in future
+        # versions.
+        self.instruction = instruction
+
+        prompt: str = self.rename_variable_question
+        # Get the disassembly lines and add them to the prompt.
+        for line in instruction.instruction_operands:
+            prompt += str(line)
 
-    def send_query(self, query: str) -> None:
+        return prompt
+
+    def rename_variable(self, response: str) -> None:
+        '''Renames the variable of the instruction saved in the Agent instance
+        to the response passed in as an argument.
+        '''
+        if self.instruction is None:
+            raise TypeError('No instruction was saved in the Agent instance.')
+        if response is None or response == '':
+            raise TypeError(f'No response was returned from OpenAI; got type {type(response)}.')
+        # Get just one word from the response.
+        response = response.split()[0]
+        # Assign the variable name to the response.
+        log.log_debug(f'Renaming variable in expression {self.instruction} to {response}.')
+        self.instruction.dest.name = response
+
+
+    def send_query(self, query: str, callback: Optional[Callable]=None) -> None:
         '''Sends a query to the engine and prints the response.'''
         query = Query(query_string=query,
                       model=self.model,
-                      max_token_count=self.get_token_count())
+                      max_token_count=self.get_token_count(),
+                      callback_function=callback)
         query.start()

src/entry.py

@@ -1,6 +1,8 @@
 from pathlib import Path
 from binaryninja import BinaryView, Function
-from binaryninja.highlevelil import HighLevelILInstruction
+from binaryninja.highlevelil import HighLevelILInstruction, HighLevelILVarInit, \
+                                    HighLevelILFunction
+from binaryninja.log import log_error
 from . agent import Agent
 
 API_KEY_PATH = Path.home() / Path('.openai/api_key.txt')
@@ -13,9 +15,31 @@ def check_function(bv: BinaryView, func: Function) -> bool:
     query: str = agent.generate_query(func)
     agent.send_query(query)
 
-def rename_expression(bv: BinaryView, instruction: HighLevelILInstruction) -> bool:
+def rename_variable(bv: BinaryView, instruction: HighLevelILInstruction) -> bool:
+
+    if not isinstance(instruction, HighLevelILVarInit):
+        log_error(f'Instruction must be of type HighLevelILVarInit, got type: ' \
+                  f'{type(instruction)}')
+        return False
+
     agent: Agent = Agent(
         bv=bv,
         path_to_api_key=API_KEY_PATH
     )
-    query: str = agent.generate_rename_expression_query(instruction)
+    query: str = agent.generate_rename_variable_query(instruction)
+    agent.send_query(query=query, callback=agent.rename_variable)
+
+# Difficult to test without a payment method added, given that the rate limits
+# are so low. This should also probably take place in a background task of its
+# own.
+# def rename_all_variables_in_function(bv: BinaryView, func: HighLevelILFunction) -> None:
+#     # Get each instruction in the High Level IL Function.
+#     for instruction in func.instructions:
+#         match instruction:
+#             # Rename the variable if it is a HighLevelILVarInit.
+#             case HighLevelILVarInit():
+#                 rename_variable(bv, instruction)
+#             # Explicit pass for all other cases.
+#             case _ :
+#                 pass
+

src/query.py

@@ -1,17 +1,20 @@
 import openai
+from typing import Callable
 from binaryninja.plugin import BackgroundTaskThread
-
+from binaryninja.highlevelil import HighLevelILInstruction
+from binaryninja.log import log_info
 
 class Query(BackgroundTaskThread):
 
     def __init__(self, query_string: str, model: str,
-                 max_token_count: int) -> None:
+                 max_token_count: int, callback_function: Callable=None) -> None:
         BackgroundTaskThread.__init__(self,
                                       initial_progress_text="",
                                       can_cancel=False)
         self.query_string: str = query_string
         self.model: str = model
         self.max_token_count: int = max_token_count
+        self.callback = callback_function
 
     def run(self) -> None:
         self.progress = "Submitting query to OpenAI."
@@ -21,5 +24,11 @@ def run(self) -> None:
             prompt=self.query_string,
             max_tokens=self.max_token_count,
         )
-        # Notify the user.
-        print(response.choices[0].text)
+        # Get the response text.
+        response: str = response.choices[0].text
+        # If there is a callback, do something with it.
+        if self.callback:
+            self.callback(response)
+        # Otherwise, assume we just want to log it.
+        else:
+            log_info(response)