51 lines
1.9 KiB
Python
51 lines
1.9 KiB
Python
import os
|
|
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
|
|
|
|
# 忽略输出torch cpu执行警告
|
|
warnings.filterwarnings('ignore', category=UserWarning)
|
|
|
|
|
|
def addr2vec(base_file_path):
|
|
# 从路径拆分文件名与路径
|
|
file_name = str(os.path.basename(base_file_path))
|
|
file_path = str(os.path.dirname(base_file_path))
|
|
|
|
# 忽略已生成的文件
|
|
if os.path.exists(os.path.join(file_path, 'final', file_name)):
|
|
return
|
|
|
|
# 如果不是路径则开始转化
|
|
if file_name:
|
|
file_json = load_json(base_file_path)
|
|
# 确保存在基础文件而不存在特征文件的情况
|
|
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:
|
|
feature_set = {}
|
|
for item in feature_json:
|
|
feature_set[item['addr']] = bb2vec(item['opcode'])
|
|
for item in file_json['acfg_list']:
|
|
bb_feature_addr_list = item['block_features']
|
|
item['block_features'] = [feature_set[key] for key in bb_feature_addr_list]
|
|
save_json(os.path.join(file_path, 'final', file_name), file_json)
|
|
else:
|
|
logger.error(f'文件{file_name}不存在特征文件')
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
logger = setup_logger('feature2json', './log/feature2json.log')
|
|
sample_type = 'benign'
|
|
json_path = os.path.join(f'./out/json/{sample_type}')
|
|
json_files = os.listdir(json_path)
|
|
with multiprocessing.Pool(processes=os.cpu_count()) as pool:
|
|
result = list(tqdm(pool.imap_unordered(addr2vec, [os.path.join(json_path, file) for file in json_files]),
|
|
total=len(json_files)))
|
|
|
|
# for json_file in json_files:
|
|
# addr2vec(json_path, json_file)
|