import concurrent.futures
import os
import r2pipe
from tqdm import tqdm


def get_fun_name_list(file_path):
    # 读取csv文件
    r2 = r2pipe.open(os.path.join(file_path), flags=['-2'])
    r2.cmd('aaa')
    r2.cmd('e arch=x86')
    function_list = r2.cmdj("aflj")
    fun_name_list = []
    for function in function_list:
        fun_name_list.append(function['name'])
    r2.quit()
    return fun_name_list


if __name__ == '__main__':
    file_path = os.path.join('/mnt/d/bishe/dataset/sample_20230130_458')
    file_list = os.listdir(file_path)
    fun_name_set = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=6) as executor:
        future_to_args = {
            executor.submit(get_fun_name_list, os.path.join(file_path, file_name)): file_name
            for file_name in file_list
        }
        for future in tqdm(concurrent.futures.as_completed(future_to_args), total=len(future_to_args)):
            fun_name_list = future.result()
            for fun_name in fun_name_list:
                if fun_name not in fun_name_set:
                    fun_name_set[fun_name] = 1
                else:
                    fun_name_set[fun_name] += 1
    print(fun_name_set)