76 lines
3.0 KiB
Python
76 lines
3.0 KiB
Python
import json
|
|
import os
|
|
|
|
import torch
|
|
from torch_geometric.data import Data
|
|
from tqdm import tqdm
|
|
|
|
from src.utils.Vocabulary import Vocab
|
|
|
|
|
|
def parse_json_list_2_pyg_object(jsonl_file: str, label: int, vocab: Vocab, save_path: str, file_type: str):
|
|
# def parse_json_list_2_pyg_object(jsonl_file: str):
|
|
train_type = ['train', 'valid', 'test']
|
|
index = 0
|
|
file_index = 0
|
|
type_index = 0
|
|
valid_flag = True
|
|
test_flag = True
|
|
file_len = len(os.listdir(jsonl_file))
|
|
|
|
for file in tqdm(os.listdir(jsonl_file)):
|
|
if index >= file_len * 0.8 and valid_flag:
|
|
type_index += 1
|
|
valid_flag = False
|
|
file_index = 0
|
|
print("make valid set")
|
|
elif index >= file_len * 0.9 and test_flag:
|
|
type_index += 1
|
|
test_flag = False
|
|
file_index = 0
|
|
print("make test set")
|
|
j = json_to_pt(file=jsonl_file + file, label=label, vocab=vocab, save_path=save_path, file_type=file_type, train_type=train_type[type_index], index=file_index)
|
|
index += 1
|
|
file_index += 1
|
|
|
|
|
|
def json_to_pt(file: str, label: int, vocab: Vocab, save_path: str, file_type: str, train_type: str, index: int):
|
|
with open(file, "r", encoding="utf-8") as item:
|
|
line = item.readline()
|
|
item = json.loads(line)
|
|
item_hash = item['hash']
|
|
acfg_list = []
|
|
for one_acfg in item['acfg_list']: # list of dict of acfg
|
|
block_features = one_acfg['block_features']
|
|
block_edges = one_acfg['block_edges']
|
|
one_acfg_data = Data(x=torch.tensor(block_features, dtype=torch.float),
|
|
edge_index=torch.tensor(block_edges, dtype=torch.long))
|
|
acfg_list.append(one_acfg_data)
|
|
|
|
item_function_names = item['function_names']
|
|
item_function_edges = item['function_edges']
|
|
|
|
local_function_name_list = item_function_names[:len(acfg_list)]
|
|
assert len(acfg_list) == len(
|
|
local_function_name_list), "The length of ACFG_List should be equal to the length of Local_Function_List"
|
|
external_function_name_list = item_function_names[len(acfg_list):]
|
|
|
|
external_function_index_list = [vocab[f_name] for f_name in external_function_name_list]
|
|
|
|
torch.save(Data(hash=item_hash, local_acfgs=acfg_list, external_list=external_function_index_list,
|
|
function_edges=item_function_edges, targets=label),
|
|
save_path + "{}_{}/{}_{}.pt".format(train_type, file_type, file_type, index))
|
|
return True
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
json_path = "./jsonl/infected_jsonl/"
|
|
train_vocab_file = "../data/processed_dataset/train_external_function_name_vocab.jsonl"
|
|
save_vocab_file = "../data/processed_dataset/DatasetJSON/"
|
|
file_type = "malware"
|
|
max_vocab_size = 10000
|
|
vocabulary = Vocab(freq_file=train_vocab_file, max_vocab_size=max_vocab_size)
|
|
parse_json_list_2_pyg_object(jsonl_file=json_path, label=1, vocab=vocabulary, save_path=save_vocab_file,
|
|
file_type=file_type)
|