import concurrent.futures
import os
import r2pipe
from tqdm import tqdm
import pandas as pd



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

        for function in function_list:
            fun_name_list.append(function['name'])
    except Exception as err:
        print(f'error at {file_path} , {err}')
    r2.quit()
    return fun_name_list

def fun_name_count():
    file_path = os.path.join('/mnt/d/bishe/dataset/sample_20230130_458')
    bengin_file_path = os.path.join('/mnt/d/bishe/dataset/train_benign')
    file_list = [os.path.join(file_path, file_name) for file_name in os.listdir(file_path)]
    file_list.extend([os.path.join(bengin_file_path, file_name) for file_name in os.listdir(bengin_file_path)])
    fun_name_set = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=12) as executor:
        future_to_args = {
            executor.submit(get_fun_name_list, 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()
            if fun_name_list:
                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
    pd.DataFrame(fun_name_set.items(), columns=['fun_name', 'count']).to_csv('./out/fun_name.csv', index=False, mode='a')

def fun_name_sort():
    fun_name_df = pd.read_csv('./out/fun_name.csv')
    # 去除fun_name中fun_name列中的局部函数
    for item in ['fcn.', 'loc.', 'main', 'entr']:
        fun_name_df = fun_name_df[fun_name_df['fun_name'].apply(lambda x: item not in x and item not in x)]
    fun_name_df = fun_name_df.sort_values(by='count', ascending=False)[:10000]
    fun_name_df.to_csv('fun_name_sort.csv', index=False)

if __name__ == '__main__':
    fun_name_count()
    fun_name_sort()