Inst2Vec/process_data/r2test/max_op_count.py
2024-04-11 16:43:57 +08:00

49 lines
1.9 KiB
Python

import r2pipe
from my_utils import setup_logger, multi_thread, THREAD_FULL
import os
from tqdm import tqdm
def get_all_from_exe(file):
# 获取基础块内的操作码序列
r2pipe_open = r2pipe.open(os.path.join(file), flags=['-2'])
try:
# 获取函数列表
r2pipe_open.cmd("aaa")
r2pipe_open.cmd('e arch=x86')
function_list = r2pipe_open.cmdj("aflj")
exe_op_count = []
for function in function_list:
function_op_count_list = []
if function['name'][:4] not in ['fcn.', 'loc.', 'main', 'entr']:
continue
block_list = r2pipe_open.cmdj("afbj @" + str(function['offset']))
for block in block_list:
# 获取基本块的反汇编指令
disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"]))
block_op_count = 0
if disasm:
print_flag = 1 if len(disasm) >= 723 else 0
for op in disasm:
if op["type"] == "invalid" or op["opcode"] == "invalid":
continue
if print_flag == 1:
print(op['disasm'])
block_op_count += 1
function_op_count_list.append(block_op_count)
exe_op_count.append(function_op_count_list)
logger.info(f"{file} {exe_op_count}")
except Exception as e:
logger.error(f"Error: get function list failed in {file} ,error {e}")
return False, file, e
r2pipe_open.quit()
return True, '', ''
if __name__ == '__main__':
logger = setup_logger('get_all_from_exe', '../../log/get_all_from_exe.log')
file = '/mnt/d/bishe/dataset/sample_benign'
file_list = os.listdir(file)
multi_thread(get_all_from_exe, ['/mnt/d/bishe/dataset/sample_benign/00125dcd81261701fcaaf84d0cb45d0e.exe'], thread_num=THREAD_FULL)