66from openai .api_resources .model import Model
77from openai .error import APIError
88
9+ from binaryninja .function import Function
910from binaryninja .lowlevelil import LowLevelILFunction
1011from binaryninja .mediumlevelil import MediumLevelILFunction
1112from binaryninja .highlevelil import HighLevelILFunction
1213from binaryninja .settings import Settings
13- from binaryninja import log
14+ from binaryninja import log , BinaryView
1415
1516from . query import Query
17+ from . c import Pseudo_C
1618
1719
1820class Agent :
1921
2022 question : str = '''
2123 This is a function that was decompiled with Binary Ninja.
22- It is in Binary Ninja's IL_FORM. What does this function do?
24+ It is in IL_FORM. What does this function do?
2325 '''
2426
2527 # A mapping of IL forms to their names.
2628 il_name : dict [type , str ] = {
2729 LowLevelILFunction : 'Low Level Intermediate Language' ,
2830 MediumLevelILFunction : 'Medium Level Intermediate Language' ,
29- HighLevelILFunction : 'High Level Intermediate Language'
31+ HighLevelILFunction : 'High Level Intermediate Language' ,
32+ Function : 'decompiled C code'
3033 }
3134
3235 def __init__ (self ,
33- function : Union [LowLevelILFunction , MediumLevelILFunction , HighLevelILFunction ],
36+ bv : BinaryView ,
37+ function : Union [Function , LowLevelILFunction ,
38+ MediumLevelILFunction , HighLevelILFunction ],
3439 path_to_api_key : Optional [Path ]= None ) -> None :
3540
3641 # Read the API key from the environment variable.
@@ -39,12 +44,16 @@ def __init__(self,
3944 # Ensure that a function type was passed in.
4045 if not isinstance (
4146 function ,
42- (LowLevelILFunction , MediumLevelILFunction , HighLevelILFunction )):
47+ (Function , LowLevelILFunction , MediumLevelILFunction ,
48+ HighLevelILFunction )):
4349 raise TypeError (f'Expected a BNIL function of type '
44- f'LowLevelILFunction, MediumLevelILFunction, or '
45- f'HighLevelILFunction, got { type (function )} .' )
50+ f'Function, LowLevelILFunction, '
51+ f'MediumLevelILFunction, or HighLevelILFunction, '
52+ f'got { type (function )} .' )
4653
54+ assert bv is not None , 'BinaryView is None. Check how you called this function.'
4755 # Set instance attributes.
56+ self .bv = bv
4857 self .function = function
4958 self .model = self .get_model ()
5059
@@ -124,14 +133,17 @@ def instruction_list(self, function: Union[LowLevelILFunction,
124133 '''Generates a list of instructions in string representation given a
125134 BNIL function.
126135 '''
136+ if isinstance (function , Function ):
137+ return Pseudo_C (self .bv , function ).get_c_source ()
127138 instructions : list [str ] = []
128139 for instruction in function .instructions :
129140 instructions .append (str (instruction ))
130141 return instructions
131142
132- def generate_query (self , function : Union [LowLevelILFunction ,
133- MediumLevelILFunction ,
134- HighLevelILFunction ]) -> str :
143+ def generate_query (self , function : Union [Function ,
144+ LowLevelILFunction ,
145+ MediumLevelILFunction ,
146+ HighLevelILFunction ]) -> str :
135147 '''Generates a query string given a BNIL function. Reads the file
136148 prompt.txt and replaces the IL form with the name of the IL form.
137149 '''
0 commit comments