Merge pull request #15 from WhatTheFuzz/feature/settings

Implement user settings

.pylintrc

@@ -103,7 +103,10 @@ disable=
     # We anticipate #3512 where it will become optional
     fixme,
     consider-alternative-union-syntax,
-    relative-beyond-top-level
+    relative-beyond-top-level,
+    # Remove import error for clients without the Binary Ninja plugin installed,
+    # as in non-commercial settings.
+    import-error
 
 
 [REPORTS]

README.md

@@ -23,9 +23,16 @@ please submit a pull request if you've tested it.
 ## API Key
 
 This requires an [API token from OpenAI][token]. The plugin checks for the API
-key in two ways (in this order).
+key in three ways (in this order).
 
-First, it checks the environment variable `OPENAI_API_KEY`, which you can set
+First, it tries to read the key from Binary Ninja's preferences. You can
+access the entry in Binary Ninja via `Edit > Preferences > Settings > OpenAI`.
+Or, use the hotkey ⌘+, and search for `OpenAI`. You should see customizable
+settings like so.
+
+![Settings](https://github.com/WhatTheFuzz/binaryninja-openai/blob/main/resources/settings.png?raw=true)
+
+Second, it checks the environment variable `OPENAI_API_KEY`, which you can set
 inside of Binary Ninja's Python console like so:
 
 ```python
@@ -42,8 +49,8 @@ mkdir ~/.openai
 echo -n "INSERT KEY HERE" > ~/.openai/api_key.txt
 ```
 
-Note that if you have both set, the plugin defaults to the environment variable.
-If your API token is invalid, you'll receive the following error:
+Note that if you have all three set, the plugin defaults to one set in Binary
+Ninja. If your API token is invalid, you'll receive the following error:
 
 ```python
 openai.error.AuthenticationError: Incorrect API key provided: <BAD KEY HERE>.
@@ -64,12 +71,14 @@ inside of the function.
 
 The output will appear in Binary Ninja's Log like so:
 
-![The output of running the plugin.](./resources/output.png)
+![The output of running the plugin.](https://github.com/WhatTheFuzz/binaryninja-openai/blob/main/resources/output.png?raw=true)
 
 ## OpenAI Model
 
 By default, the plugin uses the `text-davinci-003` model, you can tweak this
-inside of [entry.py][entry].
+inside Binary Ninja's preferences. You can access these settings as described in
+the [API Key](#api-key) section. It uses the maximum available number of tokens
+for each model, as described in [OpenAI's documentation][tokens].
 
 ## Known Issues
 
@@ -86,6 +95,7 @@ This project is licensed under the [MIT license][license].
 
 [default-plugin-dir]:https://docs.binary.ninja/guide/plugins.html
 [token]:https://beta.openai.com/account/api-keys
+[tokens]:https://beta.openai.com/docs/models/gpt-3
 [entry]:./src/entry.py
 [asyncio]:https://docs.python.org/3/library/asyncio.html
 [issue-8]:https://github.com/WhatTheFuzz/binaryninja-openai/issues/8

__init__.py

@@ -1,6 +1,10 @@
 from binaryninja import PluginCommand
+from . src.settings import OpenAISettings
 from . src.entry import check_function
 
+# Register the settings group in Binary Ninja to store the API key and model.
+OpenAISettings()
+
 PluginCommand.register_for_high_level_il_function("OpenAI\What Does this Function Do (HLIL)?",
                             "Checks OpenAI to see what this HLIL function does." \
                             "Requires an internet connection and an API key "

plugin.json

@@ -9,7 +9,7 @@
       "python3"
    ],
    "description": "Queries OpenAI's GPT3 to determine what a given function does.",
-   "longdescription": "Generates a query that asks 'What does this function do?' followed by a list of the instructions in the function. Returns the result from GPT3 and displays it to the user in the Binary Ninja console. Requires an OpenAI API key.",
+   "longdescription": "",
    "license": {
       "name": "MIT",
       "text": "Copyright 2022 Sean Deaton (@WhatTheFuzz)\n\nPermission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE."
@@ -27,6 +27,6 @@
          "openai"
       ]
 	},
-   "version": "1.0.0",
+   "version": "1.1.0",
    "minimumbinaryninjaversion": 3200
 }

src/agent.py

@@ -3,16 +3,15 @@
 from pathlib import Path
 
 import openai
-from openai.api_resources.engine import Engine
+from openai.api_resources.model import Model
 from openai.error import APIError
 
 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 .exceptions import InvalidEngineException
-
 
 class Agent:
 
@@ -30,7 +29,6 @@ class Agent:
 
     def __init__(self,
                 function: Union[LowLevelILFunction, MediumLevelILFunction, HighLevelILFunction],
-                engine: str,
                 path_to_api_key: Optional[Path]=None) -> None:
 
         # Read the API key from the environment variable.
@@ -44,33 +42,82 @@ def __init__(self,
                             f'LowLevelILFunction, MediumLevelILFunction, or '
                             f'HighLevelILFunction, got {type(function)}.')
 
-        # Get the list of available engines.
-        engines: list[Engine] = openai.Engine.list().data
-        # Ensure the user's selected engine is available.
-        if engine not in [e.id for e in engines]:
-            InvalidEngineException(f'Invalid engine: {engine}. Valid engines '
-                                   f'are: {[e.id for e in engines]}')
-
         # Set instance attributes.
         self.function = function
-        self.engine = engine
+        self.model = self.get_model()
 
     def read_api_key(self, filename: Optional[Path]=None) -> str:
-        if os.getenv('OPENAI_API_KEY'):
-            return os.getenv('OPENAI_API_KEY')
+        '''Checks for the API key in three locations.
+
+        First, it checks the openai.api_key key:value in Binary Ninja
+        preferences. This is accessed in Binary Ninja by going to Edit >
+        Preferences > Settings > OpenAI.
+        Second, it checks the OPENAI_API_KEY environment variable.
+        Finally, it checks the file specified by the filename argument.
+        Defaults to ~/.openai/api_key.txt.
+        '''
+
+        # First, check the Binary Ninja settings.
+        settings: Settings = Settings()
+        if settings.contains('openai.api_key'):
+            if key := settings.get_string('openai.api_key'):
+                return key
+
+        # If the settings don't exist, contain the key, or the key is empty,
+        # check the environment variable.
+        if key := os.getenv('OPENAI_API_KEY'):
+            return key
+
+        # Finally, if the environment variable doesn't exist, check the default
+        # file.
         if filename:
             log.log_info(f'No API key detected under the environment variable '
                          f'OPENAI_API_KEY. Reading API key from {filename}')
             try:
                 with open(filename, mode='r', encoding='ascii') as api_key_file:
                     return api_key_file.read()
-            except FileNotFoundError as error:
+            except FileNotFoundError:
                 log.log_error(f'Could not find API key file at {filename}.')
 
-        raise APIError('No API key found. Please set the environment '
-                        'variable OPENAI_API_KEY to your API key, or write '
-                        'it to ~/openai/api_key.txt.')
+        raise APIError('No API key found. Refer to the documentation to add the '
+                       'API key.')
+
+    def is_valid_model(self, model: str) -> bool:
+        '''Checks if the model is valid by querying the OpenAI API.'''
+        models: list[Model] = openai.Model.list().data
+        return model in [m.id for m in models]
 
+    def get_model(self) -> str:
+        '''Returns the model that the user has selected from Binary Ninja's
+        preferences. The default value is set by the OpenAISettings class. If
+        for some reason the user selected a model that doesn't exist, this
+        function defaults to 'text-davinci-003'.
+        '''
+        settings: Settings = Settings()
+        # Check that the key exists.
+        if settings.contains('openai.model'):
+            # Check that the key is not empty and get the user's selection.
+            if model := settings.get_string('openai.model'):
+                # Check that is a valid model by querying the OpenAI API.
+                if self.is_valid_model(model):
+                    return model
+        # Return a valid, default model.
+        assert self.is_valid_model('text-davinci-003')
+        return 'text-davinci-003'
+
+    def max_token_count(self, model: str) -> int:
+        '''Returns the maximum number of tokens that can be generated by the
+        model. Returns a default of 2,048 if the model is not found. '''
+        # TODO: This should be somewhere else, as it's also shared by Settings.
+        models: dict[str, int] = {
+            'text-davinci-003': 4_000,
+            'text-curie-001': 2_048,
+            'text-babbage-001': 2_048,
+            'text-ada-001': 2_048,
+            'code-davinci-002': 8_000,
+            'code-cushman-001': 2_048
+        }
+        return models.get(model, 2_048)
 
     def instruction_list(self, function: Union[LowLevelILFunction,
                                          MediumLevelILFunction,
@@ -101,9 +148,8 @@ def generate_query(self, function: Union[LowLevelILFunction,
     def send_query(self, query: str) -> str:
         '''Sends a query to the engine and returns the response.'''
         response: str = openai.Completion.create(
-            model=self.engine,
+            model=self.model,
             prompt=query,
-            max_tokens=2_048
+            max_tokens=self.max_token_count(self.model) - len(query),
         )
         return response.choices[0].text
-

src/entry.py

@@ -9,7 +9,6 @@
 def check_function(bv: BinaryView, func: Function) -> bool:
     agent: Agent = Agent(
         function=func,
-        engine='text-davinci-003',
         path_to_api_key=API_KEY_PATH
     )
     query: str = agent.generate_query(func)

src/exceptions.py

@@ -1,4 +1,5 @@
-from openai.error import OpenAIError
+class RegisterSettingsGroupException(Exception):
+    pass
 
-class InvalidEngineException(OpenAIError):
+class RegisterSettingsKeyException(Exception):
     pass

src/settings.py

@@ -0,0 +1,66 @@
+import json
+from binaryninja.settings import Settings
+from . exceptions import RegisterSettingsGroupException, \
+                         RegisterSettingsKeyException
+
+class OpenAISettings(Settings):
+
+    def __init__(self) -> None:
+        # Initialize the settings with the default instance ID.
+        super().__init__(instance_id='default')
+        # Register the OpenAI group.
+        if not self.register_group('openai', 'OpenAI'):
+            raise RegisterSettingsGroupException('Failed to register OpenAI '
+                                                 'settings group.')
+        # Register the setting for the API key.
+        if not self.register_api_key_settings():
+            raise RegisterSettingsKeyException('Failed to register OpenAI API '
+                                               'key settings.')
+
+        # Register the setting for the model used to query.
+        if not self.register_model_settings():
+            raise RegisterSettingsKeyException('Failed to register OpenAI '
+                                               'model settings.')
+
+    def register_api_key_settings(self) -> bool:
+        '''Register the OpenAI API key settings in Binary Ninja.'''
+        # Set the attributes of the settings. Refer to:
+        # https://api.binary.ninja/binaryninja.settings-module.html
+        properties = {
+            'title': 'OpenAI API Key',
+            'type': 'string',
+            'description': 'The user\'s OpenAI API key used to make requests '
+            'the server.'
+        }
+        return self.register_setting('openai.api_key', json.dumps(properties))
+
+    def register_model_settings(self) -> bool:
+        '''Register the OpenAI model settings in Binary Ninja.
+        Defaults to text-davinci-003.
+        '''
+        # Set the attributes of the settings. Refer to:
+        # https://api.binary.ninja/binaryninja.settings-module.html
+        properties = {
+            'title': 'OpenAI Model',
+            'type': 'string',
+            'description': 'The OpenAI model used to generate the response.',
+            # https://beta.openai.com/docs/models
+            'enum': [
+                'text-davinci-003',
+                'text-curie-001',
+                'text-babbage-001',
+                'text-babbage-002',
+                'code-davinci-002',
+                'code-cushman-001'
+            ],
+            'enumDescriptions': [
+                'Most capable GPT-3 model. Can do any task the other models can do, often with higher quality, longer output and better instruction-following. Also supports inserting completions within text.',
+                'Very capable, but faster and lower cost than Davinci.',
+                'Capable of straightforward tasks, very fast, and lower cost.',
+                'Capable of very simple tasks, usually the fastest model in the GPT-3 series, and lowest cost.',
+                'Most capable Codex model. Particularly good at translating natural language to code. In addition to completing code, also supports inserting completions within code.',
+                'Almost as capable as Davinci Codex, but slightly faster. This speed advantage may make it preferable for real-time applications.'
+            ],
+            'default': 'text-davinci-003'
+        }
+        return self.register_setting('openai.model', json.dumps(properties))