bert与radare拆分

exe2json.py

@@ -13,6 +13,8 @@ import os
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 ret_trap_opcode_family = ["ret", "hlt", "int3", "ud2"]
 
+sample_type = 'malware'
+
 
 def extract_opcode(disasm_text):
     """
@@ -39,13 +41,21 @@ def calc_sha256(file_path):
         return sha256
 
 
-def get_graph_cfg_r2pipe(r2pipe_open, file_path):
+def get_graph_cfg_r2pipe(r2pipe_open, file_path, feature_out):
     # CFG提取
     acfg_item = []
+    acfg_feature_item = []
     try:
         # 获取函数列表
         function_list = r2pipe_open.cmdj("aflj")
+        # debug
+        # for function in tqdm(function_list, total=len(function_list)):
         for function in function_list:
+            # print(function['offset'])
+            # debug
+            # if function['offset'] != 5368850704:
+            #     continue
+
             node_list = []
             edge_list = []
             temp_edge_list = []
@@ -54,8 +64,6 @@ def get_graph_cfg_r2pipe(r2pipe_open, file_path):
             block_list = r2pipe_open.cmdj("afbj @" + str(function['offset']))
             # 获取基本块数量
             block_number = len(block_list)
-
-
             for block in block_list:
                 # 基础块内的语句
                 block_addr = block["addr"]
@@ -64,11 +72,21 @@ def get_graph_cfg_r2pipe(r2pipe_open, file_path):
                 node_list.append(block["addr"])
                 # 获取基本块的反汇编指令
                 disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"]))
+                # debug
+                # if len(disasm) != 49144:
+                #     continue
                 if disasm:
+                    last_op = ''
+                    if len(disasm) > 200:
+                        logger.warning(f"基础块指令长度异常，{file_path}，函数名称{function['name']}基础块地址{block['addr']},长度{len(disasm)}")
                     for op_index, op in enumerate(disasm):
+                        op_disasm = extract_opcode(op["disasm"])
+                        # 防止大量重复的语句造成内存溢出
+                        if op_disasm == last_op:
+                            continue
+                        last_op = op_disasm
                         # 提取操作码并转换为bert模型输入格式
                         block_Statement.append(extract_opcode(op["disasm"]))
-
                         # 处理跳转码并构建cfg
                         if 'jump' in op:
                             if op['jump'] == 0:
@@ -109,16 +127,30 @@ def get_graph_cfg_r2pipe(r2pipe_open, file_path):
                                 if op not in ret_trap_opcode_family or op["type"] not in ["ret", "trap"]:
                                     temp_edge_list.append([block_addr, op["offset"] + op["size"]])
 
+                # debugger
+                # print(len(disasm))
+                # print(len(block_Statement))
+                # print(block_Statement)
+
+
+                """
+                速度过慢
+                """
                 # bert模型转化特征
-                block_feature_list = bb2vec(block_Statement)
+                # block_feature_list = bb2vec(block_Statement)
                 # block_feature_list = []
+
+                # 暂时将bb地址作为特征  后续将操作码集中转化为特征
+                block_feature_list = block_addr
+                acfg_feature_item.append({'addr':block_addr, 'opcode':block_Statement})
+
             # 过滤不存在的边
             for temp_edge in temp_edge_list:
                 if temp_edge[0] in node_list and temp_edge[1] in node_list:
                     edge_list.append(temp_edge)
             # 单独错误信息日志
-            if block_number == 0 or len(block_feature_list) == 0:
-                logger.warning(f"二进制可执行文件解析出错，出错文件：{file_path}，出错函数地址：{function['offset']},基础块个数{block_number},基础块特征{block_feature_list}")
+            if block_number == 0:
+                logger.warning(f"二进制可执行文件解析出错，出错文件：{file_path}，出错函数地址：{function['offset']},基础块个数{block_number}")
             # cfg构建
             acfg = {
                 'block_number': block_number,
@@ -126,6 +158,8 @@ def get_graph_cfg_r2pipe(r2pipe_open, file_path):
                 'block_features': block_feature_list
             }
             acfg_item.append(acfg)
+            feature_out.write(json.dumps(acfg_feature_item))
+
         return True, "二进制可执行文件解析成功", acfg_item
     except Exception as e:
         return False, e, None
@@ -184,21 +218,26 @@ def get_r2pipe(file_path):
 
 def init_logging():
     # 初始化日志
-    log_file = "./log/exe2json.log"
+    log_file = f"./log/exe2json_{sample_type}.log"
     return setup_logger('exe2json', log_file)
 
 
 def exe_to_json(file_path):
-    output_path = "./out/json/malware"
-
-
+    output_path = f"./out/json/{sample_type}"
+    file_fingerprint = calc_sha256(file_path)
+    if os.path.exists(os.path.join(output_path, file_fingerprint + '.jsonl')):
+        logger.info(f"二进制可执行文件已解析，文件名{file_path}")
+        return True
+    logger.info(f"开始解析，文件名{file_path}")
     # 获取r2pipe并解析文件  解析完即释放r2
     r2 = get_r2pipe(file_path)
     fcg_Operation_flag, fcg_Operation_message, function_num, function_fcg_edge_list, function_names = get_graph_fcg_r2pipe(r2)
-    cfg_Operation_flag, cfg_Operation_message, cfg_item = get_graph_cfg_r2pipe(r2,file_path)
+    with open(os.path.join(output_path, 'feature' ,file_fingerprint + '.jsonl'), 'w') as feature_out:
+        cfg_Operation_flag, cfg_Operation_message, cfg_item = get_graph_cfg_r2pipe(r2,file_path, feature_out)
+        feature_out.close()
     r2.quit()
     # 文件json构建
-    file_fingerprint = calc_sha256(file_path)
+
     if fcg_Operation_flag and cfg_Operation_flag:
         json_obj = {
             'hash': file_fingerprint,
@@ -221,12 +260,14 @@ def exe_to_json(file_path):
     with open(os.path.join(output_path, file_fingerprint + '.jsonl'), 'w') as out:
         out.write(result)
         out.close()
+    logger.info(f"解析完成，文件名{file_path}")
     return True
 
 if __name__ == '__main__':
     logger = init_logging()
-    sample_file_path = "/mnt/d/bishe/dataset/sample_malware"
+
+    sample_file_path = f"/mnt/d/bishe/dataset/sample_{sample_type}"
     sample_file_list = os.listdir(sample_file_path)
-    multi_thread(exe_to_json, [os.path.join(sample_file_path, file_name) for file_name in sample_file_list])
+    multi_thread(exe_to_json, [os.path.join(sample_file_path, file_name) for file_name in sample_file_list], thread_num=THREAD_FULL)
     # test_file_path = '/mnt/d/bishe/exe2json/sample/VirusShare_0a3b625380161cf92c4bb10135326bb5'
     # exe_to_json(test_file_path)

external_func/funNameGet.py

@@ -17,21 +17,22 @@ def get_fun_name_list(file_path):
 
         for function in function_list:
             fun_name_list.append(function['name'])
+            print(function['name'])
     except Exception as err:
         print(f'error at {file_path} , {err}')
     r2.quit()
     return fun_name_list
 
 def fun_name_count():
-    file_path = os.path.join('/mnt/d/bishe/dataset/sample_20230130_458')
-    bengin_file_path = os.path.join('/mnt/d/bishe/dataset/train_benign')
+    file_path = os.path.join('/mnt/d/bishe/dataset/sample_malware')
+    bengin_file_path = os.path.join('/mnt/d/bishe/dataset/sample_benign')
     file_list = [os.path.join(file_path, file_name) for file_name in os.listdir(file_path)]
     file_list.extend([os.path.join(bengin_file_path, file_name) for file_name in os.listdir(bengin_file_path)])
     fun_name_set = {}
     with concurrent.futures.ThreadPoolExecutor(max_workers=12) as executor:
         future_to_args = {
             executor.submit(get_fun_name_list, file_name): file_name
-            for file_name in file_list
+            for file_name in file_list[:1]
         }
         for future in tqdm(concurrent.futures.as_completed(future_to_args), total=len(future_to_args)):
             fun_name_list = future.result()
@@ -41,16 +42,16 @@ def fun_name_count():
                         fun_name_set[fun_name] = 1
                     else:
                         fun_name_set[fun_name] += 1
-    pd.DataFrame(fun_name_set.items(), columns=['fun_name', 'count']).to_csv('./out/fun_name.csv', index=False, mode='a')
+    # pd.DataFrame(fun_name_set.items(), columns=['fun_name', 'count']).to_csv('../out/external/fun_name.csv', index=False, mode='a')
 
 def fun_name_sort():
-    fun_name_df = pd.read_csv('./out/fun_name.csv')
+    fun_name_df = pd.read_csv('../out/external/fun_name.csv')
     # 去除fun_name中fun_name列中的局部函数
     for item in ['fcn.', 'loc.', 'main', 'entr']:
         fun_name_df = fun_name_df[fun_name_df['fun_name'].apply(lambda x: item not in x and item not in x)]
     fun_name_df = fun_name_df.sort_values(by='count', ascending=False)[:10000]
-    fun_name_df.to_csv('fun_name_sort.csv', index=False)
+    fun_name_df.to_csv('../out/external/fun_name_sort.csv', index=False)
 
 if __name__ == '__main__':
     fun_name_count()
-    fun_name_sort()
+    # fun_name_sort()

log_utils.py

@@ -1,45 +0,0 @@
-import logging
-import os
-
-
-def setup_logger(name, log_file, level=logging.INFO):
-    """Function setup as many loggers as you want"""
-    if not os.path.exists(os.path.dirname(log_file)):
-        os.makedirs(os.path.dirname(log_file))
-
-    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
-
-    handler = logging.FileHandler(log_file)
-    handler.setFormatter(formatter)
-
-    # 控制台是否输出日志信息
-    # stream_handler = logging.StreamHandler()
-    # stream_handler.setFormatter(formatter)
-
-    logger = logging.getLogger(name)
-    logger.setLevel(level)
-    logger.addHandler(handler)
-    # 控制台
-    # logger.addHandler(stream_handler)
-
-    # 刷新原有log文件
-
-    if os.path.exists(log_file):
-        open(log_file, 'w').close()
-
-    return logger
-
-
-"""
-用法示例
-"""
-
-
-def main():
-    log_file = "app.log"
-    logger = setup_logger(__name__, log_file)
-
-    logger.info("Application started.")
-    logger.debug("Debug message.")
-    logger.warning("Warning message.")
-    logger.error("Error occurred.")

ngram/OpcodeGet.py

@@ -56,19 +56,8 @@ def get_graph_r2pipe(file_type, file_name):
         r2pipe_open.cmd('e arch=x86')
         function_list = r2pipe_open.cmdj("aflj")
         for function in function_list:
-
-            # 外部函数测试
-            # if function['name'] == 'sub.TNe_U':
-            #     print(function)
-            #     block_list = r2pipe_open.cmdj("afbj @" + str(function['offset']))
-            #     for block in block_list:
-            #         # print(block)
-            #         # 获取基本块的反汇编指令
-            #         disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"]))
-            #         if disasm:
-            #             for op in disasm:
-            #                 print(extract_opcode(op["opcode"]))
-
+            if function['name'][:4] not in ['fcn.', 'loc.']:
+                continue
             block_list = r2pipe_open.cmdj("afbj @" + str(function['offset']))
             block_opcode_Sequence = []
             for block in block_list:
@@ -92,30 +81,27 @@ def get_graph_r2pipe(file_type, file_name):
 if __name__ == '__main__':
     file_type = 'malware'
     logger = setup_logger('logger', f'./log/opcode_{file_type}.log')
-    file_path = os.path.join('/mnt/d/bishe/dataset/sample_20230130_458')
+    file_path = os.path.join(f'/mnt/d/bishe/dataset/sample_{file_type}')
     print(f"max works {os.cpu_count()}")
-    file_list = os.listdir(file_path)[:10000]
+    file_list = os.listdir(file_path)[:1000]
     done_file_num = 0
     done_list = [['class', 'sub-class', 'size', 'corpus']]
-    process_bar = tqdm(desc=f'Processing {file_type}...', leave=True, total=len(file_list))
-    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:  # 调整线程池大小
+    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()/2) as executor:  # 调整线程池大小
         future_to_args = {
             executor.submit(get_graph_r2pipe, file_type, file_name): file_name for file_name in file_list
         }
-        for future in concurrent.futures.as_completed(future_to_args):
+        for future in tqdm(concurrent.futures.as_completed(future_to_args), total=len(future_to_args), desc=f'Processing {file_type}...',):
             try:
                 tmp = future.result()
                 done_list.extend(tmp if len(tmp) > 0 else [])
                 if len(done_list) > 100000:
-                    csv_write(f'output_{file_type}.csv', done_list)
+                    csv_write(f'output_{file_type}_test.csv', done_list)
                     done_file_num += 1
                     done_list.clear()
             except Exception as e:
                 logger.error(f"Error: {e}")
                 print(f"Error: {e}")
-            finally:
-                process_bar.update(1)
         else:
-            csv_write(f'output_{file_type}.csv', done_list)
+            csv_write(f'output_{file_type}_test.csv', done_list)
 
 

ngram/ngram.py

@@ -93,11 +93,17 @@ def process_csv_file(csvfile, ngram_type, file_percent_filter, frequency_filter)
     print(f"start load csv file:{os.path.basename(csvfile)}")
     dataframe = pd.read_csv(csvfile, encoding="utf8")
     print(f"end load")
+    print(f"computer {idx + 1} ngram for {ngram_type}")
+    # reader = pd.read_csv(csvfile, chunksize=10000)
+    print(f"init file {csvfile} done")
     ngram_list = defaultdict(int)
     filtered_ngram_list = defaultdict(int)
-    process_bar = tqdm(total=len(dataframe['corpus'].values), desc=f'Computing {ngram_type}-gram on files')
     with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor:  # 调整线程池大小
         future_to_args = {
+           # executor.submit(filter_N_grams, values['corpus'].values, idx + 1, file_percent_filter, frequency_filter):
+              #  values for values in reader
+
+
             executor.submit(filter_N_grams, dataframe['corpus'].values[start: start + 10000],
                             idx + 1, file_percent_filter, frequency_filter): start for start in
             range(0, len(dataframe['corpus'].values), 10000)
@@ -128,7 +134,7 @@ def process_csv_file(csvfile, ngram_type, file_percent_filter, frequency_filter)
 def build_csv(ngram_list, filter_list, maxgrams, file_type):
     ngramDicList = []
     csv_file_header = ['ngram', 'count']
-    csv_file = os.path.join('./out', f'{file_type}-{maxgrams}-gram.csv')
+    csv_file = os.path.join('../out', f'{file_type}-{maxgrams}-gram.csv')
     for index in tqdm(range(len(ngram_list)), desc=f'Building {maxgrams}-gram csv'):
         ngramDicList.append({
             'ngram': ngram_list[index],