Merge pull request #19 from WhatTheFuzz/feature/C

Feature/c

Author: WhatTheFuzz

README.md

@@ -4,7 +4,8 @@
 
 Integrates OpenAI's GPT3 with Binary Ninja via a plugin. Creates a query asking
 "What does this function do?" followed by the instructions in the High Level IL
-function. Returns the response to the user in Binary Ninja's console.
+function or the decompiled pseudo-C. Returns the response to the user in Binary
+Ninja's console.
 
 ## Installation
 

__init__.py

@@ -11,3 +11,10 @@
                             "saved under the environment variable "
                             "OPENAI_API_KEY or modify the path in entry.py.",
                             check_function)
+
+PluginCommand.register_for_function("OpenAI\What Does this Function Do (Pseudo-C)?",
+                            "Checks OpenAI to see what this pseudo-C function does." \
+                            "Requires an internet connection and an API key "
+                            "saved under the environment variable "
+                            "OPENAI_API_KEY or modify the path in entry.py.",
+                            check_function)

plugin.json

@@ -19,14 +19,14 @@
       "Linux"
    ],
    "installinstructions": {
-      "Darwin": "pip3 install --user openai. Under your OpenAI account, create an API key and set it to the environment variable OPENAI_API_KEY inside Binary Ninja's Python console.",
-      "Linux": "pip3 install --user openai. Under your OpenAI account, create an API key and set it to the environment variable OPENAI_API_KEY inside Binary Ninja's Python console."
+      "Darwin": "`pip3 install --user openai`.\n Add your OpenAI API key to the OpenAI preferences in Binary Ninja.",
+      "Linux": "`pip3 install --user openai`.\n Add your OpenAI API key to the OpenAI preferences in Binary Ninja."
    },
    "dependencies": {
       "pip": [
          "openai"
       ]
 	},
-   "version": "1.3.0",
+   "version": "1.4.0",
    "minimumbinaryninjaversion": 3200
 }

src/agent.py

@@ -6,31 +6,36 @@
 from openai.api_resources.model import Model
 from openai.error import APIError
 
+from binaryninja.function import Function
 from binaryninja.lowlevelil import LowLevelILFunction
 from binaryninja.mediumlevelil import MediumLevelILFunction
 from binaryninja.highlevelil import HighLevelILFunction
 from binaryninja.settings import Settings
-from binaryninja import log
+from binaryninja import log, BinaryView
 
 from . query import Query
+from . c import Pseudo_C
 
 
 class Agent:
 
     question: str = '''
     This is a function that was decompiled with Binary Ninja.
-    It is in Binary Ninja's IL_FORM. What does this function do?
+    It is in IL_FORM. What does this function do?
     '''
 
     # A mapping of IL forms to their names.
     il_name: dict[type, str] = {
         LowLevelILFunction: 'Low Level Intermediate Language',
         MediumLevelILFunction: 'Medium Level Intermediate Language',
-        HighLevelILFunction: 'High Level Intermediate Language'
+        HighLevelILFunction: 'High Level Intermediate Language',
+        Function: 'decompiled C code'
     }
 
     def __init__(self,
-                function: Union[LowLevelILFunction, MediumLevelILFunction, HighLevelILFunction],
+                bv: BinaryView,
+                function: Union[Function, LowLevelILFunction,
+                                MediumLevelILFunction, HighLevelILFunction],
                 path_to_api_key: Optional[Path]=None) -> None:
 
         # Read the API key from the environment variable.
@@ -39,12 +44,16 @@ def __init__(self,
         # Ensure that a function type was passed in.
         if not isinstance(
                 function,
-            (LowLevelILFunction, MediumLevelILFunction, HighLevelILFunction)):
+            (Function, LowLevelILFunction, MediumLevelILFunction,
+                                                HighLevelILFunction)):
             raise TypeError(f'Expected a BNIL function of type '
-                            f'LowLevelILFunction, MediumLevelILFunction, or '
-                            f'HighLevelILFunction, got {type(function)}.')
+                            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()
 
@@ -124,14 +133,17 @@ def instruction_list(self, function: Union[LowLevelILFunction,
         '''Generates a list of instructions in string representation given a
         BNIL function.
         '''
+        if isinstance(function, Function):
+            return Pseudo_C(self.bv, function).get_c_source()
         instructions: list[str] = []
         for instruction in function.instructions:
             instructions.append(str(instruction))
         return instructions
 
-    def generate_query(self, function: Union[LowLevelILFunction,
-                                       MediumLevelILFunction,
-                                       HighLevelILFunction]) -> str:
+    def generate_query(self, function: Union[Function,
+                                            LowLevelILFunction,
+                                            MediumLevelILFunction,
+                                            HighLevelILFunction]) -> str:
         '''Generates a query string given a BNIL function. Reads the file
         prompt.txt and replaces the IL form with the name of the IL form.
         '''

src/entry.py

@@ -4,13 +4,11 @@
 
 API_KEY_PATH = Path.home() / Path('.openai/api_key.txt')
 
-# We don't use the bv argument, but it gets passed in by the PluginCommand.
-# pylint: disable=unused-argument
 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)
-    response: str = agent.send_query(query)
-    print(response)
+    agent.send_query(query)