MalGraph/samples/funCount.py
2024-01-10 10:32:38 +08:00

39 lines
1.5 KiB
Python

import json
import os
from itertools import islice
import heapq
from tqdm import tqdm
if __name__ == '__main__':
mal_file_name = './jsonl/infected_jsonl/'
ben_file_name = './jsonl/refind_jsonl/'
fun_name_dict = {}
for file in tqdm(os.listdir(mal_file_name)):
with open(mal_file_name + file, 'r') as item:
item = json.loads(item.readline())
item_fun_list = item['function_names']
for fun_name in item_fun_list:
if fun_name != 'start' and fun_name != 'start_0' and 'sub_' not in fun_name:
if fun_name_dict.get(fun_name) is not None:
fun_name_dict[fun_name] += 1
else:
fun_name_dict[fun_name] = 1
for file in tqdm(os.listdir(ben_file_name)):
with open(ben_file_name + file, 'r') as item:
item = json.loads(item.readline())
item_fun_list = item['function_names']
for fun_name in item_fun_list:
if fun_name != 'start' and fun_name != 'start_0' and 'sub_' not in fun_name:
if fun_name_dict.get(fun_name) is not None:
fun_name_dict[fun_name] += 1
else:
fun_name_dict[fun_name] = 1
with open('./res.jsonl', 'w') as file:
largest_10000_items = heapq.nlargest(10000, fun_name_dict.items(), key=lambda item: item[1])
for key, value in largest_10000_items:
temp = {"f_name": key, "count": value}
file.write(json.dumps(temp) + '\n')