Remove function parameter from agent.

The function is passed in to generate the query; it does not need to be known at initialization time. Additionally, having that parameter prevents us from registering PluginCommands that don't work on Function (HLIL or otherwise) objects. In this case, we want to register an instruction, not a function.

Author: WhatTheFuzz

__init__.py

@@ -1,6 +1,7 @@
 from binaryninja import PluginCommand
 from . src.settings import OpenAISettings
 from . src.entry import check_function
+from . src.entry import rename_expression
 
 # Register the settings group in Binary Ninja to store the API key and model.
 OpenAISettings()
@@ -18,3 +19,11 @@
  "saved under the environment variable "
  "OPENAI_API_KEY or modify the path in entry.py.",
  check_function)
+
+PluginCommand.register_for_high_level_il_instruction("OpenAI\Rename Variable",
+ "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)

src/agent.py

@@ -34,27 +34,14 @@ class Agent:
 
  def __init__(self,
  bv: BinaryView,
- function: Union[Function, LowLevelILFunction,
- MediumLevelILFunction, HighLevelILFunction],
  path_to_api_key: Optional[Path]=None) -> None:
 
  # Read the API key from the environment variable.
  openai.api_key = self.read_api_key(path_to_api_key)
 
- # Ensure that a function type was passed in.
- if not isinstance(
- function,
- (Function, LowLevelILFunction, MediumLevelILFunction,
- HighLevelILFunction)):
- raise TypeError(f'Expected a BNIL function of type '
- f'Function, LowLevelILFunction, '
- f'MediumLevelILFunction, or HighLevelILFunction, '
- f'got {type(function)}.')
-
  assert bv is not None, 'BinaryView is None. Check how you called this function.'
  # Set instance attributes.
  self.bv = bv
- self.function = function
  self.model = self.get_model()
 
  def read_api_key(self, filename: Optional[Path]=None) -> str:
@@ -133,6 +120,15 @@ def instruction_list(self, function: Union[LowLevelILFunction,
  '''Generates a list of instructions in string representation given a
  BNIL function.
  '''
+
+ # Ensure that a function type was passed in.
+ if not isinstance(function, (Function, LowLevelILFunction,
+ MediumLevelILFunction, HighLevelILFunction)):
+ raise TypeError(f'Expected a BNIL function of type '
+ f'Function, LowLevelILFunction, '
+ f'MediumLevelILFunction, or HighLevelILFunction, '
+ f'got {type(function)}.')
+
  if isinstance(function, Function):
  return Pseudo_C(self.bv, function).get_c_source()
  instructions: list[str] = []

src/entry.py

@@ -1,14 +1,17 @@
 from pathlib import Path
 from binaryninja import BinaryView, Function
+from binaryninja.highlevelil import HighLevelILInstruction
 from . agent import Agent
 
 API_KEY_PATH = Path.home() / Path('.openai/api_key.txt')
 
 def check_function(bv: BinaryView, func: Function) -> bool:
  agent: Agent = Agent(
  bv=bv,
- function=func,
  path_to_api_key=API_KEY_PATH
  )
  query: str = agent.generate_query(func)
  agent.send_query(query)
+
+def rename_expression(bv: BinaryView, instruction: HighLevelILInstruction) -> bool:
+ pass