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 # 禁用分词器多线程 os.environ["TOKENIZERS_PARALLELISM"] = "false" 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, 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 = [] 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: 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: 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"二进制可执行文件解析警告，跳转指令识别出错，指令{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"]]) # 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 = 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) feature_out.write(json.dumps(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 node_bytes = "" node_opcode = "" for op in pdf["ops"]: if op["type"] == "invalid": continue node_bytes += op["bytes"] opcode = extract_opcode(op["disasm"]) node_opcode += opcode + " " 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) sub_function_name_list = ('sym.','sub','imp') func_name_list = [func_name for func_name in func_name_list if func_name.startswith(sub_function_name_list)] return True, "二进制可执行文件解析成功", function_num, edge_list, func_name_list except Exception as e: return False, e, None, None, None def get_r2pipe(file_path): # 初始化r2pipe try: r2 = r2pipe.open(file_path, flags=['-2']) r2.cmd("aaa") r2.cmd('e arch=x86') return r2 except Exception as e: return 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 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) 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构建 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 False # json写入 result = json.dumps(json_obj,ensure_ascii=False, default=lambda x: float(x) if isinstance(x, np.float32) else x) os.makedirs(output_path, exist_ok=True) 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 = 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], thread_num=THREAD_FULL) # test_file_path = '/mnt/d/bishe/exe2json/sample/VirusShare_0a3b625380161cf92c4bb10135326bb5' # exe_to_json(test_file_path)