Skip to content

Commit 2f3692a

Browse files
robertmarsalrobertmarsal
committed
feat(PLU-212): add support for RevEng user agent
1 parent 670cfe4 commit 2f3692a

12 files changed

Lines changed: 50 additions & 33 deletions

File tree

‎reai_toolkit/features/ai_decompiler/ai_decompiler.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,7 @@ def start_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
9090
raise Exception("Analysis not found. Please choose one using the 'Attach to existing' feature.")
9191
function_id = get_function_id_by_addr_util(bv, function.start, self.config)
9292

93-
with revengai.ApiClient(self.config.api_config) as api_client:
93+
with self.config.create_api_client() as api_client:
9494
api_instance = revengai.FunctionsAIDecompilationApi(api_client)
9595
api_response_status = api_instance.get_ai_decompilation_task_status(function_id)
9696
log_info(f"RevEng.AI | AI Decompilation task created for function at 0x{function.start:x}")
@@ -106,7 +106,7 @@ def start_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
106106

107107
if poll_status.lower() == "uninitialised":
108108
try:
109-
with revengai.ApiClient(self.config.api_config) as api_client:
109+
with self.config.create_api_client() as api_client:
110110
api_instance = revengai.FunctionsAIDecompilationApi(api_client)
111111
api_response_status = api_instance.create_ai_decompilation_task(function_id)
112112
log_info(f"RevEng.AI | AI Decompilation task created for function at 0x{function.start:x}")
@@ -126,7 +126,7 @@ def start_ai_decompiler(self, bv: BinaryView, options: Dict) -> None:
126126

127127
if poll_status.lower() == "completed":
128128
log_info(f"RevEng.AI | AI Decompilation for function at 0x{function.start:x} is completed")
129-
with revengai.ApiClient(self.config.api_config) as api_client:
129+
with self.config.create_api_client() as api_client:
130130
api_instance = revengai.FunctionsAIDecompilationApi(api_client)
131131
api_response_status = api_instance.get_ai_decompilation_task_result(function_id, summarise=True, generate_inline_comments=True)
132132

‎reai_toolkit/features/auto_unstrip/auto_unstrip.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,7 @@ def auto_unstrip(self, bv: BinaryView):
109109
if self.cancelled.is_set():
110110
return False, "Operation cancelled"
111111

112-
with revengai.ApiClient(self.config.api_config) as api_client:
112+
with self.config.create_api_client() as api_client:
113113
api_instance = revengai.FunctionsCoreApi(api_client)
114114
api_response = api_instance.auto_unstrip(analysis_id, auto_unstrip_request)
115115

‎reai_toolkit/features/choose_source/choose_source.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ def get_analysis(self, bv: BinaryView):
4444

4545
sha256 = get_sha256(bv.file.filename)
4646

47-
with revengai.ApiClient(self.config.api_config) as api_client:
47+
with self.config.create_api_client() as api_client:
4848
api_instance = revengai.SearchApi(api_client)
4949
api_response = api_instance.search_binaries(partial_sha256 = sha256, user_files_only = True)
5050
results = api_response.data.results

‎reai_toolkit/features/configuration/config.py‎

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ def __init__(self):
1515
self.sha256 = ""
1616
self.api_config = None
1717
self.is_configured = False
18+
self.user_agent = "Binary Ninja/RevEng.AI_Plugin"
19+
self.reveng_header = "X-RevEng-Application"
1820
self._load_config()
1921

2022

@@ -46,7 +48,6 @@ def _load_config(self):
4648
self.api_key = settings.get_string("revengai.api_key", None)
4749
self.portal_url = settings.get_string("revengai.portal_url", None)
4850
self.init_api_config()
49-
#re_conf["user_agent"] = "RevEng.AI BinaryNinja Plugin"
5051

5152

5253
def save_config(self) -> bool:
@@ -75,8 +76,24 @@ def init_api_config(self):
7576
self.api_config = revengai.Configuration(api_key={'APIKey': self.api_key}, host=self.host)
7677
except Exception as e:
7778
log_error(f"RevEng.AI | Failed to initialize API client: {str(e)}")
78-
79-
79+
80+
def create_api_client(self):
81+
"""
82+
Create a properly configured RevEng.AI API client with custom headers and user agent.
83+
84+
Returns:
85+
revengai.ApiClient: Configured API client instance
86+
87+
Usage:
88+
with self.config.create_api_client() as api_client:
89+
api_instance = revengai.SomeApi(api_client)
90+
# ... use api_instance
91+
"""
92+
api_client = revengai.ApiClient(self.api_config, self.reveng_header, self.user_agent)
93+
api_client.user_agent = self.user_agent
94+
return api_client
95+
96+
8097
def clear_config(self):
8198
self.api_key = ""
8299
self.host = ""
@@ -88,7 +105,7 @@ def clear_config(self):
88105
def check_auth(self):
89106
try:
90107
self.init_api_config()
91-
with revengai.ApiClient(self.api_config) as api_client:
108+
with self.create_api_client() as api_client:
92109
api_instance = revengai.AuthenticationUsersApi(api_client)
93110
api_response = api_instance.get_requester_user_info()
94111
log_info(f"RevEng.AI | Welcome {api_response.data.username}!")
@@ -149,7 +166,7 @@ def init_config(self, bv: BinaryView):
149166
AnalysisSyncService(self).sync_analysis_data(analysis_id=self.analysis_id, bv=bv)
150167
else:
151168
log_info(f"RevEng.AI | Binary not found in saved configurations, searching in RevEng.AI...")
152-
with revengai.ApiClient(self.api_config) as api_client:
169+
with self.create_api_client() as api_client:
153170
api_instance = revengai.SearchApi(api_client)
154171
api_response = api_instance.search_binaries(partial_sha256 = self.sha256, user_files_only = True)
155172
if not len(api_response.data.results):

‎reai_toolkit/features/match_current_function/match_current_function.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ def parse_confidence(item):
5959
if self.cancelled.is_set():
6060
return False, "Operation cancelled"
6161

62-
with revengai.ApiClient(self.config.api_config) as api_client:
62+
with self.config.create_api_client() as api_client:
6363
analysis_core_instance = revengai.AnalysesResultsMetadataApi(api_client)
6464
analyzed_functions = analysis_core_instance.get_functions_list(analysis_id)
6565
analyzed_functions = analyzed_functions.to_dict()["data"]["functions"]
@@ -103,7 +103,7 @@ def parse_confidence(item):
103103
matched_count = 0
104104
while True:
105105
time.sleep(3)
106-
with revengai.ApiClient(self.config.api_config) as api_client:
106+
with self.config.create_api_client() as api_client:
107107
analysis_core_instance = revengai.FunctionsCoreApi(api_client)
108108
functions_by_distance = analysis_core_instance.batch_function_matching( schema_ann_model)
109109

‎reai_toolkit/features/match_functions/match_functions.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def parse_confidence(item):
4949
if self.cancelled.is_set():
5050
return False, "Operation cancelled"
5151

52-
with revengai.ApiClient(self.config.api_config) as api_client:
52+
with self.config.create_api_client() as api_client:
5353
analysis_core_instance = revengai.AnalysesResultsMetadataApi(api_client)
5454
analyzed_functions = analysis_core_instance.get_functions_list(analysis_id)
5555
analyzed_functions = analyzed_functions.to_dict()["data"]["functions"]
@@ -119,7 +119,7 @@ def parse_confidence(item):
119119
matched_count = 0
120120
while True:
121121
time.sleep(3)
122-
with revengai.ApiClient(self.config.api_config) as api_client:
122+
with self.config.create_api_client() as api_client:
123123
analysis_core_instance = revengai.FunctionsCoreApi(api_client)
124124
functions_by_distance = analysis_core_instance.analysis_function_matching(analysis_id, schema_ann_model)
125125

‎reai_toolkit/features/upload/upload.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ def __init__(self, config):
1111

1212
def get_models(self, bv: BinaryView):
1313
try:
14-
with revengai.ApiClient(self.config.api_config) as api_client:
14+
with self.config.create_api_client() as api_client:
1515
api_instance = revengai.ModelsApi(api_client)
1616
api_response = api_instance.get_models()
1717
models = api_response.data.models
@@ -24,7 +24,7 @@ def get_models(self, bv: BinaryView):
2424
def upload_file(self, file_path: str, file_type: revengai.UploadFileType):
2525
try:
2626
file = Path(file_path).read_bytes()
27-
with revengai.ApiClient(self.config.api_config) as api_client:
27+
with self.config.create_api_client() as api_client:
2828
api_instance = revengai.AnalysesCoreApi(api_client)
2929
api_response = api_instance.upload_file(file_type, file, force_overwrite=True)
3030
log_info(f"RevEng.AI | File uploaded successfully.")
@@ -71,7 +71,7 @@ def upload_binary(self, bv: BinaryView, options: dict):
7171
})
7272
log_info(f"RevEng.AI | Collected {len(symbols['function_boundaries'])} functions from Binary Ninja!")
7373

74-
with revengai.ApiClient(configuration=self.config.api_config) as api_client:
74+
with self.config.create_api_client() as api_client:
7575

7676
analyses_client = revengai.AnalysesCoreApi(api_client)
7777

@@ -95,7 +95,7 @@ def upload_binary(self, bv: BinaryView, options: dict):
9595

9696
log_info(f"RevEng.AI | Analysis started successfully. Analysis ID: {analysis_result.data.analysis_id}, Binary ID: {analysis_result.data.binary_id}")
9797

98-
PeriodicChecker(self.config).start_checking(bv, analysis_result.data.analysis_id, analysis_result.data.binary_id, self.config.set_current_info, self.config.api_config)
98+
PeriodicChecker(self.config).start_checking(bv, analysis_result.data.analysis_id, analysis_result.data.binary_id, self.config.set_current_info, self.config)
9999

100100
return True, "Analysis started successfully."
101101

‎reai_toolkit/features/view_function_in_portal/view_function_in_portal.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ def view_function_in_portal(self, bv: BinaryView, options: Dict) -> None:
2626
if not analysis_id:
2727
raise Exception("Analysis not found. Please choose one using the 'Attach to existing' feature.")
2828

29-
with revengai.ApiClient(self.config.api_config) as api_client:
29+
with self.config.create_api_client() as api_client:
3030
api_instance = revengai.AnalysesResultsMetadataApi(api_client)
3131
api_response = api_instance.get_functions_list(analysis_id)
3232
analyzed_functions = api_response.data.functions

‎reai_toolkit/utils/core/binary_ninja.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
def _rename_in_portal(config: revengai.Configuration, function_id:int, new_name:str, new_mangled_name:str):
1010
try:
11-
with revengai.ApiClient(config.api_config) as api_client:
11+
with config.create_api_client() as api_client:
1212
api_instance = revengai.FunctionsRenamingHistoryApi(api_client)
1313
api_instance.rename_function_id(
1414
function_id=function_id,
@@ -71,7 +71,7 @@ def get_function_by_addr(bv: BinaryView, addr: int) -> Function:
7171
return functions_containing[0]
7272

7373
def get_function_id_by_addr(bv: BinaryView, addr: int, config):
74-
with revengai.ApiClient(config.api_config) as api_client:
74+
with config.create_api_client() as api_client:
7575
api_instance = revengai.AnalysesResultsMetadataApi(api_client)
7676
analysis_id = config.get_analysis_id(bv)
7777
api_response = api_instance.get_functions_list(analysis_id)

‎reai_toolkit/utils/features/matching.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ def _search_collection(self, query: Dict[str, Any] = {}):
8888
try:
8989
output = []
9090
log_info("RevEng.AI | Searching for collections")
91-
with revengai.ApiClient(self.config.api_config) as api_client:
91+
with self.config.create_api_client() as api_client:
9292
api_instance = revengai.SearchApi(api_client)
9393
api_response = api_instance.search_collections(
9494
page = 1,
@@ -116,7 +116,7 @@ def _search_binaries(self, query: Dict[str, Any] = {}):
116116
output = []
117117
try:
118118
log_info("RevEng.AI | Searching for binaries")
119-
with revengai.ApiClient(self.config.api_config) as api_client:
119+
with self.config.create_api_client() as api_client:
120120
api_instance = revengai.SearchApi(api_client)
121121
api_response = api_instance.search_binaries(
122122
page = 1,
@@ -184,7 +184,7 @@ def search_collections(self, query: Dict[str, Any] = {}) -> None:
184184
page = 1
185185
while True:
186186
log_info(f"RevEng.AI | Searching for collections on page {page}")
187-
with revengai.ApiClient(self.config.api_config) as api_client:
187+
with self.config.create_api_client() as api_client:
188188
api_instance = revengai.SearchApi(api_client)
189189
api_response = api_instance.search_collections(
190190
page = page,
@@ -220,7 +220,7 @@ def search_binaries(self, query: Dict[str, Any] = {}) -> None:
220220
page = 1
221221
while True:
222222
log_info(f"RevEng.AI | Searching for binaries on page {page}")
223-
with revengai.ApiClient(self.config.api_config) as api_client:
223+
with self.config.create_api_client() as api_client:
224224
api_instance = revengai.SearchApi(api_client)
225225
api_response = api_instance.search_binaries(
226226
page = page,
@@ -262,7 +262,7 @@ def _process_data_type_batch(self, chunk: List[Dict], chunk_index: int) -> List[
262262
if self.cancelled.is_set():
263263
return []
264264

265-
with revengai.ApiClient(self.config.api_config) as api_client:
265+
with self.config.create_api_client() as api_client:
266266
api_instance = revengai.FunctionsDataTypesApi(api_client)
267267
function_data_types_params = revengai.FunctionDataTypesParams.from_dict({"function_ids": function_ids})
268268
api_response = api_instance.generate_function_data_types_for_functions(function_data_types_params)
@@ -275,7 +275,7 @@ def _process_data_type_batch(self, chunk: List[Dict], chunk_index: int) -> List[
275275
while True:
276276
if self.cancelled.is_set():
277277
return []
278-
with revengai.ApiClient(self.config.api_config) as api_client:
278+
with self.config.create_api_client() as api_client:
279279
api_instance = revengai.FunctionsDataTypesApi(api_client)
280280
api_response = api_instance.list_function_data_types_for_functions(function_ids=function_ids).to_dict()
281281
log_info("The response of FunctionsDataTypesApi->list_function_data_types_for_functions:\n")

0 commit comments

Comments
 (0)