58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
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'])
|
|
print(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_malware')
|
|
bengin_file_path = os.path.join('/mnt/d/bishe/dataset/sample_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[:1]
|
|
}
|
|
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/external/fun_name.csv', index=False, mode='a')
|
|
|
|
def fun_name_sort():
|
|
fun_name_df = pd.read_csv('../out/external/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('../out/external/fun_name_sort.csv', index=False)
|
|
|
|
if __name__ == '__main__':
|
|
fun_name_count()
|
|
# fun_name_sort()
|