import random import r2pipe import hashlib from my_utils import * import json # 基础块抽取 # from bert.obtain_inst_vec import bb2vec from tqdm import tqdm import numpy as np import os import concurrent.futures import multiprocessing ret_trap_opcode_family = ["ret", "hlt", "int3", "ud2"] sample_type = 'malware' def extract_opcode(disasm_text): """ 从反汇编文本中提取操作码和操作数 正则表达式用于匹配操作码和操作数,考虑到操作数可能包含空格和逗号 将操作码与操作数转化为bert模型输入 """ op_list = disasm_text.split(' ') res = [] for item in op_list: item = item.strip().replace(',', '') if '[' in item: res.append('[') res.append(item.replace('[', '').replace(']', '')) if ']' in item: res.append(']') return ' '.join(res) def calc_sha256(file_path): with open(file_path, 'rb') as f: bytes = f.read() sha256obj = hashlib.sha256(bytes) sha256 = sha256obj.hexdigest() return sha256 def get_graph_cfg_r2pipe(r2pipe_open, file_path, output_path, file_name): # 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 = [] block_feature_list = [] # 基本快块列表 block_list = r2pipe_open.cmdj("afbj @" + str(function['offset'])) # 获取基本块数量 block_number = len(block_list) for block in block_list: # 基础块内的语句 block_addr = block["addr"] block_Statement = [] node_list.append(block["addr"]) # 获取基本块的反汇编指令 disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"])) # debug # if len(disasm) != 49144: # continue if disasm: block_len = len(disasm) for op_index, op in enumerate(disasm): # 提取操作码并转换为bert模型输入格式 op_disasm = extract_opcode(op["disasm"]) # 如果单个基础块的长度大于20且操作码重复,则跳过 if block_len > 20 and op_disasm in block_Statement: continue block_Statement.append(op_disasm) # 处理跳转码并构建cfg if 'jump' in op: if op['jump'] == 0: if op_index != len(disasm) - 1: node_list.append(disasm[op_index + 1]['offset']) elif op['type'] == 'jmp': temp_edge_list.append([block["addr"], op['jump']]) if op_index != len(disasm) - 1: node_list.append(disasm[op_index + 1]['offset']) elif op['type'] == 'cjmp': temp_edge_list.append([block["addr"], op['jump']]) if op_index == len(disasm) - 1: temp_edge_list.append([block_addr, op['jump']]) else: temp_edge_list.append([block_addr, disasm[op_index + 1]["offset"]]) node_list.append(disasm[op_index + 1]["offset"]) elif op['type'] == 'call': temp_edge_list.append([block_addr, op["jump"]]) temp_edge_list.append([op["jump"], block_addr]) if op_index == len(disasm) - 1: temp_edge_list.append([block_addr, op["offset"] + op["size"]]) else: logger.warning( f"二进制可执行文件解析警告,跳转指令识别出错,文件{file_path},指令{op}") # 操作码不存在跳转指令 else: if op_index != len(disasm) - 1: # 当前指令不是基础块的最后一条指令 if op in ret_trap_opcode_family and op["type"] in ["ret", "trap"]: node_list.append(disasm[op_index + 1]["offset"]) else: # 当前指令是基础块的最后一条指令 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"]]) if block_len > 20: logger.warning(f"二进制可执行文件解析警告,基础块长度大于20,文件{file_path},基础块地址{block_addr},操作码长度{block_len}->{len(block_Statement)}") # debugger # print(len(disasm)) # print(len(block_Statement)) # print(block_Statement) """ 速度过慢 """ # bert模型转化特征 # block_feature_list = bb2vec(block_Statement) # block_feature_list = [] # 暂时将bb地址作为特征 后续将操作码集中转化为特征 block_feature_list.append(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: logger.warning( f"二进制可执行文件解析出错,出错文件:{file_path},出错函数地址:{function['offset']},基础块个数{block_number}") # cfg构建 acfg = { 'block_number': block_number, 'block_edges': [[d[0] for d in edge_list], [d[1] for d in edge_list]], 'block_features': block_feature_list } acfg_item.append(acfg) save_json(os.path.join(output_path, 'feature', file_name + '.jsonl'), acfg_feature_item) return True, "二进制可执行文件解析成功", acfg_item except Exception as e: return False, e, None def get_graph_fcg_r2pipe(r2pipe_open): # FCG提取 try: node_list = [] func_name_list = [] edge_list = [] temp_edge_list = [] function_list = r2pipe_open.cmdj("aflj") function_num = len(function_list) for function in function_list: func_name_list.append(function["name"]) r2pipe_open.cmd(f's ' + str(function["offset"])) pdf = r2pipe_open.cmdj('pdfj') if pdf is None: continue for op in pdf["ops"]: if op["type"] == "invalid": continue if "jump" in op and op["jump"] != 0: temp_edge_list.append([function["offset"], op["jump"]]) node_list.append(function["offset"]) # 完成 FCG 构建后, 检查并清理不存在的出边 for temp_edge in temp_edge_list: if temp_edge[1] in node_list: edge_list.append(temp_edge) func_name_list = [func_name for func_name in func_name_list] return True, "二进制可执行文件解析成功", function_num, edge_list, func_name_list except Exception as e: return False, e, None, None, None def init_logging(): # 初始化日志 log_file = f"./log/exe2json_{sample_type}.log" return setup_logger('exe2json', log_file) def exe_to_json(file_path): 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 # logger.info(f"开始解析,文件名{file_path}") # 获取r2pipe并解析文件 解析完即释放r2 r2 = r2pipe.open(file_path, flags=['-2']) r2.cmd("aaa") 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, feature_out) cfg_Operation_flag, cfg_Operation_message, cfg_item = get_graph_cfg_r2pipe(r2, file_path, output_path, file_fingerprint) r2.quit() # 文件json构建 if fcg_Operation_flag and cfg_Operation_flag: json_obj = { 'hash': file_fingerprint, 'function_number': function_num, 'function_edges': [[int(d[0]) for d in function_fcg_edge_list], [int(d[1]) for d in function_fcg_edge_list]], 'acfg_list': cfg_item, 'function_names': function_names } else: logger.error(f"二进制可执行文件解析失败 文件名{file_path}") if not fcg_Operation_flag: logger.error(f"fcg错误:{fcg_Operation_message}") if not cfg_Operation_flag: logger.error(f"cfg错误:{cfg_Operation_message}") return # json写入 os.makedirs(output_path, exist_ok=True) save_json(os.path.join(output_path, file_fingerprint + '.jsonl'), json_obj) # logger.info(f"解析完成,文件名{file_path}") return if __name__ == '__main__': logger = init_logging() sample_file_path = f"/mnt/d/bishe/dataset/sample_{sample_type}" sample_file_list = os.listdir(sample_file_path) print(f"max worker {os.cpu_count()}") with multiprocessing.Pool(processes=os.cpu_count()) as pool: result = list(tqdm(pool.imap_unordered(exe_to_json, [os.path.join(sample_file_path, file_name) for file_name in sample_file_list[:10000]]), total=10000)) # with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor: # futures_to_args = { # executor.submit(exe_to_json, os.path.join(sample_file_path, file_name)): file_name for file_name in # sample_file_list # } # for future in tqdm(concurrent.futures.as_completed(futures_to_args), total=len(futures_to_args)): # pass # test_file_path = '/mnt/d/bishe/exe2json/sample/VirusShare_0a3b625380161cf92c4bb10135326bb5' # exe_to_json(test_file_path)