判断空指令

json_feature2json.py

@@ -1,11 +1,11 @@
+import concurrent.futures
 import os
-from my_utils import save_json, load_json, setup_logger, multi_thread_order, THREAD_HALF, THREAD_FULL, continuation_json
+from my_utils import save_json, load_json, setup_logger
 from bert.obtain_inst_vec import bb2vec
 import multiprocessing
 from tqdm import tqdm
 import warnings
 from datetime import datetime
-import concurrent.futures
 
 warnings.filterwarnings("ignore")
 
@@ -17,6 +17,8 @@ def addr2vec(base_file_path, index):
 
     # 如果不是路径则开始转化
     if file_name:
+        # 无操作码标志位
+        none_opcode_flag = False
         # 忽略已生成的文件
         if os.path.exists(os.path.join(file_path, 'final', file_name)):
             return
@@ -25,38 +27,50 @@ def addr2vec(base_file_path, index):
         feature_json = load_json(os.path.join(file_path, 'feature', file_name)) if os.path.exists(
             os.path.join(file_path, 'feature', file_name)) else None
         if feature_json is not None:
-
+            # 如果出现无操作码的情况，直接跳过文件
-            # 对于长度过长的文件先不处理
+            for item in feature_json:
-            if len(feature_json) > 10000:
+                if len(item['opcode']) == 0:
-                data = {
+                    logger.error(f"基础块无操作码 {file_name}，地址{item['addr']}")
-                    'file_name': file_name,
+                    none_opcode_flag = True
-                    'feature_len': len(feature_json)
+            if none_opcode_flag:
-                }
-                continuation_json(os.path.join(f'./out/json/too_long_{sample_type}.json'), data)
                 return
+            # 对于长度过长的文件先不处理
+            # if len(feature_json) > 10000:
+            #     data = {
+            #         'file_name': file_name,
+            #         'feature_len': len(feature_json)
+            #     }
+            #     continuation_json(os.path.join(f'./out/json/too_long_{sample_type}.json'), data)
+            #     return
 
             # 多线程预测bert
             feature_set = {}
-            # with multiprocessing.Pool(processes=os.cpu_count()) as pool:
+            with multiprocessing.Pool(processes=4) as pool:
-            #     try:
-            #         results = list(tqdm(pool.imap_unordered(bb2vec, [item for item in feature_json]),
-            #                             total=len(feature_json),
-            #                             desc=f'{file_name} Progress:{index}/{json_files_len} ',
-            #                             leave=True,
-            #                             dynamic_ncols=True))
-            #         for result in results:
-            #             feature_set[result[0]] = result[1]
-            #     except Exception as e:
-            #         logger.error(f"bert 解析出错 {file_name}，{e}")
-            # debug
                 try:
-                for index, feature in tqdm(enumerate(feature_json), total=len(feature_json)):
+                    results = list(tqdm(pool.imap_unordered(bb2vec, [item for item in feature_json]),
-                    addr, feature = bb2vec(feature)
+                                        total=len(feature_json),
-                    feature_set[addr] = feature
+                                        desc=f'{file_name} Progress:{index}/{json_files_len} ',
+                                        leave=False,
+                                        dynamic_ncols=True,
+                                        position=1))
+                    for result in results:
+                        if result[0]:
+                            feature_set[result[1]] = result[2]
+                        else:
+                            logger.error(f"bert解析出错 {file_name}，地址{result[1]},操作码{result[2]}，报错{result[3]}")
+                            return
                 except Exception as e:
-                print(index)
+                    logger.error(f"多线程解析出错：{file_name}，报错{e}")
-                print(e)
+                    return
-                print(feature['opcode'])
+            # debug
+            # try:
+            #     for index, feature in tqdm(enumerate(feature_json), total=len(feature_json)):
+            #         addr, feature = bb2vec(feature)
+            #         feature_set[addr] = feature
+            # except Exception as e:
+            #     print(index)
+            #     print(e)
+            #     print(feature['opcode'])
 
             try:
                 for item in file_json['acfg_list']:
@@ -73,12 +87,12 @@ def addr2vec(base_file_path, index):
 
 
 if __name__ == '__main__':
-    logger = setup_logger('feature2json', './log/feature2json.log')
+    logger = setup_logger('feature2json', './log/feature2json.log', reset=True)
     sample_type = 'benign'
     # json_path = os.path.join(f'./out/json/{sample_type}')
     json_path = os.path.join(f'./out/json/{sample_type}')
-    # json_files = os.listdir(json_path)
+    json_files = os.listdir(json_path)
-    json_files = ['1710ae16c54ac149f353ba58e752ba7069f88326e6b71107598283bd0fffcbd6.jsonl']
+    # json_files = ['1710ae16c54ac149f353ba58e752ba7069f88326e6b71107598283bd0fffcbd6.jsonl']
     json_files_len = len(json_files)
     now = datetime.now()
     formatted_now = now.strftime("%Y-%m-%d %H:%M:%S")
@@ -88,6 +102,25 @@ if __name__ == '__main__':
     #                        total=len(json_files)))
     # multi_thread_order(addr2vec, [os.path.join(json_path, file) for file in json_files if
     #                               os.path.isfile(os.path.join(json_path, file))], thread_num=THREAD_FULL)
-    for index, json_file in enumerate(json_files):
+    # with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
+    #     tqdm_position = 1
+    #     try:
+    #         future_args = {
+    #             executor.submit(addr2vec, os.path.join(json_path, file), index, tqdm_position)
+    #             for index, file in enumerate(json_files)
+    #         }
+    #         for future in tqdm(concurrent.futures.as_completed(future_args),
+    #                            total=len(json_files),
+    #                            desc='Total:',
+    #                            position=0
+    #                            ):
+    #             tqdm_position += 1
+    #     except Exception as e:
+    #         print(e)
+
+    for index, json_file in tqdm(enumerate(json_files),
+                                 total=len(json_files),
+                                 desc='Total:',
+                                 position=0):
         if os.path.isfile(os.path.join(json_path, json_file)):
             addr2vec(os.path.join(json_path, json_file), index)