|
9 | 9 | from binaryninja.lowlevelil import LowLevelILFunction |
10 | 10 | from binaryninja.mediumlevelil import MediumLevelILFunction |
11 | 11 | from binaryninja.highlevelil import HighLevelILFunction |
| 12 | +from binaryninja.settings import Settings |
12 | 13 | from binaryninja import log |
13 | 14 |
|
14 | 15 | from .exceptions import InvalidEngineException |
@@ -56,20 +57,40 @@ def __init__(self, |
56 | 57 | self.engine = engine |
57 | 58 |
|
58 | 59 | def read_api_key(self, filename: Optional[Path]=None) -> str: |
59 | | - if os.getenv('OPENAI_API_KEY'): |
60 | | - return os.getenv('OPENAI_API_KEY') |
| 60 | + '''Checks for the API key in three locations. |
| 61 | +
|
| 62 | + First, it checks the openai.api_key key:value in Binary Ninja |
| 63 | + preferences. This is accessed in Binary Ninja by going to Edit > |
| 64 | + Preferences > Settings > OpenAI. |
| 65 | + Second, it checks the OPENAI_API_KEY environment variable. |
| 66 | + Finally, it checks the file specified by the filename argument. |
| 67 | + Defaults to ~/.openai/api_key.txt. |
| 68 | + ''' |
| 69 | + |
| 70 | + # First, check the Binary Ninja settings. |
| 71 | + settings: Settings = Settings() |
| 72 | + if settings.contains('openai.api_key'): |
| 73 | + if key := settings.get_string('openai.api_key'): |
| 74 | + return key |
| 75 | + |
| 76 | + # If the settings don't exist, contain the key, or the key is empty, |
| 77 | + # check the environment variable. |
| 78 | + if key := os.getenv('OPENAI_API_KEY'): |
| 79 | + return key |
| 80 | + |
| 81 | + # Finally, if the environment variable doesn't exist, check the default |
| 82 | + # file. |
61 | 83 | if filename: |
62 | 84 | log.log_info(f'No API key detected under the environment variable ' |
63 | 85 | f'OPENAI_API_KEY. Reading API key from {filename}') |
64 | 86 | try: |
65 | 87 | with open(filename, mode='r', encoding='ascii') as api_key_file: |
66 | 88 | return api_key_file.read() |
67 | | - except FileNotFoundError as error: |
| 89 | + except FileNotFoundError: |
68 | 90 | log.log_error(f'Could not find API key file at {filename}.') |
69 | 91 |
|
70 | | - raise APIError('No API key found. Please set the environment ' |
71 | | - 'variable OPENAI_API_KEY to your API key, or write ' |
72 | | - 'it to ~/openai/api_key.txt.') |
| 92 | + raise APIError('No API key found. Refer to the documentation to add the ' |
| 93 | + 'API key.') |
73 | 94 |
|
74 | 95 |
|
75 | 96 | def instruction_list(self, function: Union[LowLevelILFunction, |
|
0 commit comments