122 lines
4.6 KiB
Python
122 lines
4.6 KiB
Python
import concurrent.futures
|
|
import os
|
|
import re
|
|
from log_utils import setup_logger
|
|
from tqdm import tqdm
|
|
|
|
import r2pipe
|
|
import pandas as pd
|
|
|
|
|
|
csv_lock = 0
|
|
|
|
|
|
def Opcode_to_csv(opcode_list, file_type):
|
|
|
|
csv_write(f'output_{file_type}.csv', opcode_list)
|
|
logger.info(f"done {done_file_num} files")
|
|
|
|
|
|
|
|
def csv_write(file_name, data: list):
|
|
"""write data to csv"""
|
|
logger.info("*======================start write==================================*")
|
|
df = pd.DataFrame(data)
|
|
chunksize = 1000
|
|
for i in range(0, len(df), chunksize):
|
|
df.iloc[i:i + chunksize].to_csv(f'./out/{file_name}', mode='a', header=False, index=False)
|
|
logger.info(f"done rows {len(df)}")
|
|
logger.info("*=================write to csv success==============================*")
|
|
return True
|
|
|
|
|
|
def extract_opcode(disasm_text):
|
|
"""
|
|
从反汇编文本中提取操作码和操作数
|
|
正则表达式用于匹配操作码和操作数，考虑到操作数可能包含空格和逗号
|
|
"""
|
|
match = re.search(r"^\s*(\S+)(?:\s+(.*))?$", disasm_text)
|
|
if match:
|
|
opcode = match.group(1)
|
|
# operands_str = match.group(2) if match.group(2) is not None else ""
|
|
# split_pattern = re.compile(r",(?![^\[]*\])") # 用于切分操作数的正则表达式
|
|
# operands = split_pattern.split(operands_str)
|
|
# return opcode, [op.strip() for op in operands if op.strip()]
|
|
return opcode
|
|
return ""
|
|
|
|
|
|
def get_graph_r2pipe(file_type, file_name):
|
|
# 获取基础块内的操作码序列
|
|
r2pipe_open = r2pipe.open(os.path.join(file_path, file_name), flags=['-2'])
|
|
opcode_Sequence = []
|
|
try:
|
|
# 获取函数列表
|
|
r2pipe_open.cmd("aaa")
|
|
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"]))
|
|
|
|
block_list = r2pipe_open.cmdj("afbj @" + str(function['offset']))
|
|
block_opcode_Sequence = []
|
|
for block in block_list:
|
|
# print(block)
|
|
# 获取基本块的反汇编指令
|
|
disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"]))
|
|
if disasm:
|
|
for op in disasm:
|
|
if op["type"] == "invalid" or op["opcode"] == "invalid":
|
|
continue
|
|
block_opcode_Sequence.append(extract_opcode(op["opcode"]))
|
|
opcode_Sequence.append(
|
|
[file_type, file_type, len(block_opcode_Sequence), ' '.join(block_opcode_Sequence)])
|
|
except Exception as e:
|
|
logger.error(f"Error: get function list failed in {file_name}")
|
|
print(f"Error: get function list failed in {file_name} ,error info {e}")
|
|
r2pipe_open.quit()
|
|
return opcode_Sequence
|
|
|
|
|
|
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')
|
|
print(f"max works {os.cpu_count()}")
|
|
file_list = os.listdir(file_path)[:10000]
|
|
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: # 调整线程池大小
|
|
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):
|
|
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)
|
|
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)
|
|
|
|
|