From 636ec90a1cd627e95df0a415ef26c0393c3c2ccc Mon Sep 17 00:00:00 2001 From: TinyCaviar <379645931@qq.com> Date: Mon, 7 Aug 2023 20:48:21 +0800 Subject: [PATCH] backup --- .../convert_pkl_to_json.py | 116 +++++++++++------- Genius3/raw-feature-extractor/ida_batch.py | 18 ++- Genius3/raw-feature-extractor/read_idaFILE.py | 2 +- Genius3/raw-feature-extractor/test.py | 18 ++- 4 files changed, 106 insertions(+), 48 deletions(-) diff --git a/Genius3/raw-feature-extractor/convert_pkl_to_json.py b/Genius3/raw-feature-extractor/convert_pkl_to_json.py index 7eeceb4..c07301c 100644 --- a/Genius3/raw-feature-extractor/convert_pkl_to_json.py +++ b/Genius3/raw-feature-extractor/convert_pkl_to_json.py @@ -3,58 +3,90 @@ import pickle as pk import re import json import os +from tqdm import tqdm -def convert(): - # for workflow in range(0, 20): - workflow = 0 +def convert(start, end): + for workflow in range(start, end): + # workflow = 0 cfg_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_cfg".format(workflow) output_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_json".format(workflow) - dot_path = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_dot".format(workflow) + dot_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_dot".format(workflow) + log_path = "D:\\hkn\\infected\\datasets\\logging\\convert_log{}.log".format(workflow) + process_log_path = "D:\\hkn\\infected\\datasets\\logging\\convert_process_log{}.log".format(workflow) - for cfg in os.listdir(cfg_dir): - name = cfg[:-4] # 纯文件名，不带后缀 - cfg_file = open(os.path.join(cfg_dir, name + '.ida'), 'r') - data = pk.load(cfg_file) - cfg_file.close() + with open(log_path, 'a+') as log, open(process_log_path, 'a+') as process_log: + logged = log.readline() + if logged == '': + log_index = 0 + else: + log_index = int(logged) - # 打开dot文件获取fcg - raw_function_edges = [] - with open(os.path.join(dot_path, name + '.dot'), 'r') as dot: - for line in dot: - if '->' in line: - raw_function_edges.append(re.findall(r'\b\d+\b', line)) + for index, cfg in enumerate(tqdm(os.listdir(cfg_dir))): + if index < log_index: + continue - # 为当前pe文件创建json对象 - json_obj = { - 'hash': data.binary_name[11:], - 'function_number': data.raw_graph_list.__len__(), - 'function_edges': [[d[0] for d in raw_function_edges], [d[1] for d in raw_function_edges]], - 'acfg_list': [], - 'function_names': [] - } - # 读取pkl文件，一个acfg由一个函数分解而来 - for acfg in data.raw_graph_list: - # 这里2是因为Genius框架提取特征时将后代数量放在2 - offspring = [d.get('v')[2] for d in acfg.g.node.values()] - # 将后代数量的特征放入bb_features中 - for i, f in enumerate(acfg.bb_features): - f.append(offspring[i]) + name = cfg[:-4] # 纯文件名，不带后缀 + cfg_file = open(os.path.join(cfg_dir, name + '.ida'), 'r') + try: + data = pk.load(cfg_file) + except EOFError: + process_log.write("index {}, {} process failed. EOFError occurred.\n".format(index, cfg)) + continue + except ValueError: + process_log.write("index {}, {} process failed. ValueError occurred.\n".format(index, cfg)) + continue - acfg_item = { - 'block_number': acfg.g.__len__(), - 'block_edges': [[d[0] for d in acfg.g.edges], [d[1] for d in acfg.g.edges]], - 'block_features': acfg.bb_features - } + cfg_file.close() - json_obj['acfg_list'].append(acfg_item) - json_obj['function_names'].append(acfg.funcname) + dot_file_path = os.path.join(dot_dir, name + '.dot') + if not os.path.exists(dot_file_path): + process_log.write("index {}, {} process failed. dot file not exists.\n".format(index, cfg)) + else: + # 打开dot文件获取fcg + raw_function_edges = [] + with open(dot_file_path, 'r') as dot: + for line in dot: + if '->' in line: + raw_function_edges.append(re.findall(r'\b\d+\b', line)) - # 将结果写入json本地文件 - result = json.dumps(json_obj) - with open(os.path.join(output_dir, name + '.jsonl'), 'w') as out: - out.write(result) + # 为当前pe文件创建json对象 + json_obj = { + 'hash': data.binary_name[11:], + 'function_number': data.raw_graph_list.__len__(), + 'function_edges': [[d[0] for d in raw_function_edges], [d[1] for d in raw_function_edges]], + 'acfg_list': [], + 'function_names': [] + } + # 读取pkl文件，一个acfg由一个函数分解而来 + for acfg in data.raw_graph_list: + # 这里2是因为Genius框架提取特征时将后代数量放在2 + offspring = [d.get('v')[2] for d in acfg.g.node.values()] + # 将后代数量的特征放入bb_features中 + for i, offs in enumerate(offspring): + acfg.bb_features[i].append(offs) + + acfg_item = { + 'block_number': acfg.g.__len__(), + 'block_edges': [[d[0] for d in acfg.g.edges], [d[1] for d in acfg.g.edges]], + 'block_features': acfg.bb_features + } + + json_obj['acfg_list'].append(acfg_item) + json_obj['function_names'].append(acfg.funcname) + + # 将结果写入json本地文件 + result = json.dumps(json_obj, ensure_ascii=False) + + with open(os.path.join(output_dir, name + '.jsonl'), 'w') as out: + out.write(result) + + log.truncate(0) + log.seek(0) + log.write(str(index)) + log.flush() + process_log.write("index {}, {} process done.\n".format(index, cfg)) if __name__ == '__main__': - convert() + convert(20, 35) diff --git a/Genius3/raw-feature-extractor/ida_batch.py b/Genius3/raw-feature-extractor/ida_batch.py index 1d5ef6f..136d299 100644 --- a/Genius3/raw-feature-extractor/ida_batch.py +++ b/Genius3/raw-feature-extractor/ida_batch.py @@ -6,6 +6,8 @@ from tqdm import tqdm import time # 单个pe文件处理超时/s +# 多次处理，一批数据中只有少量文件会超时 +# 所有数据处理完成后可以对这些数据再进行一次更长超时时间的处理，若仍然超时则放弃 TIMEOUT = 60 @@ -14,7 +16,7 @@ def call_preprocess(cmd_line): def batch_mode(): - for workflow in range(0, 1): + for workflow in range(1, 20): # workflow = 0 pe_dir = 'D:\\hkn\\infected\\datasets\\virusshare_infected{}'.format(workflow) # for test @@ -36,7 +38,8 @@ def batch_mode(): # for test # cmd_line = r'idaq64 -c -A -S"D:\hkn\project_folder\Gencoding3\Genius3\raw-feature-extractor\preprocessing_ida.py {}" -oF:\iout {}'.format( # workflow, os.path.join(pe_dir, pe)) - cmd_line = r'idaq64 -c -A -S"D:\hkn\project_folder\Gencoding3\Genius3\raw-feature-extractor\preprocessing_ida.py {}" -oF:\iout {}'.format(workflow, os.path.join(pe_dir, pe)) + cmd_line = r'idaq64 -c -A -S"D:\hkn\project_folder\Gencoding3\Genius3\raw-feature-extractor\preprocessing_ida.py {}" -oF:\iout {}'.format( + workflow, os.path.join(pe_dir, pe)) p = multiprocessing.Process(target=call_preprocess, args=[cmd_line]) p.start() @@ -51,7 +54,8 @@ def batch_mode(): if flag_kill: subprocess.call('taskkill /im idaq64.exe /f') - process_log.write("index {}, {} in workflow {} stuck, process terminated.\n".format(index, pe, workflow)) + process_log.write( + "index {}, {} in workflow {} stuck, process terminated.\n".format(index, pe, workflow)) else: # 正常运行结束 log.truncate(0) @@ -59,6 +63,14 @@ def batch_mode(): log.write(str(index)) log.flush() process_log.write("index {}, {} process done.\n".format(index, pe)) + # 一次workflow结束后将所有副产物删除 + delete_output() + + +def delete_output(): + out_dir = 'F:\\iout' + os.rmdir(out_dir) + os.mkdir(out_dir) # 注意：该py文件必须放在IDA的根目录下，且必须使用cmd命令执行，否则无法链接到python库 diff --git a/Genius3/raw-feature-extractor/read_idaFILE.py b/Genius3/raw-feature-extractor/read_idaFILE.py index 8355240..463b53d 100644 --- a/Genius3/raw-feature-extractor/read_idaFILE.py +++ b/Genius3/raw-feature-extractor/read_idaFILE.py @@ -15,7 +15,7 @@ def print_obj(obj): # sub_10F20 308 反编译代码有字符串，但是这个特征提取里没有字符串 constant，可能是间接引用的，不识别。看了下所有函数的特征，几乎都没有字符串常量，可能都是写在别的地方然后引用的。 # sub_166C4 393 if __name__ == '__main__': - testpath = "D:\\hkn\\infected\\datasets\\virusshare_infected0_cfg\\VirusShare_cd53c6637ca75ac5fc1cbe6d2ced41a1.ida" + testpath = "D:\\hkn\\infected\\datasets\\virusshare_infected11_cfg\\VirusShare_5c088a2a6e0391b7c6ab22e4648eab3a.ida" fr = open(testpath, 'r') data = pickle.load(fr) #一个二进制文件的acfgs fr.close() diff --git a/Genius3/raw-feature-extractor/test.py b/Genius3/raw-feature-extractor/test.py index 3c6cb16..4a5c124 100644 --- a/Genius3/raw-feature-extractor/test.py +++ b/Genius3/raw-feature-extractor/test.py @@ -63,8 +63,22 @@ def clock(): subprocess.call('taskkill /im idaq64.exe /f') +def delete_error(): + for workflow in range(0, 35): + convert_log_path = "D:\\hkn\\infected\\datasets\\logging\\convert_process_log{}.log".format(workflow) + json_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_json".format(workflow) + + with open(convert_log_path, 'r') as log: + for line in log: + if 'Error occurred' in line: + name = line[line.find(',') + 2: line.find('.')] + '.jsonl' + # print(os.path.join(json_dir, name)) + if os.path.exists(os.path.join(json_dir, name)): + os.remove(os.path.join(json_dir, name)) + + if __name__ == '__main__': # gen_dir() # change_max_item_lines() - subprocess.call('taskkill /im idaq64.exe /f') - + # subprocess.call('taskkill /im idaq64.exe /f') + delete_error()