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)