209 lines
7.3 KiB
Python
209 lines
7.3 KiB
Python
import os
|
||
import time
|
||
import torch
|
||
from torch.utils.data import DataLoader, Dataset
|
||
from pathlib import Path
|
||
from .datatype import Tokens, Function, Instruction
|
||
from .model import ASM2VEC
|
||
from tqdm import tqdm
|
||
class AsmDataset(Dataset):
|
||
def __init__(self, x, y):
|
||
self.x = x
|
||
self.y = y
|
||
def __len__(self):
|
||
return len(self.x)
|
||
def __getitem__(self, index):
|
||
return self.x[index], self.y[index]
|
||
import csv
|
||
def csv_read(malware_num=500,benign_num=500,max_nodes=20000,min_nodes=10):
|
||
|
||
malware_cfg_list=[]
|
||
benign_cfg_list=[]
|
||
with open('../CFG_data/malware_msg.csv', 'r', encoding='utf-8') as f:
|
||
#经下述操作后，reader成为了一个可以迭代行的文件
|
||
reader = csv.reader(f)
|
||
#先拿出csv文件的首行（一般是基本名称说明的行），此时指针指向下一行
|
||
header = next(reader)
|
||
print(header)
|
||
|
||
for row in reader:
|
||
file_name=row[0]
|
||
nodes_num=row[1]
|
||
edgs_num=row[2]
|
||
if int(nodes_num) <= max_nodes and int(nodes_num)>=min_nodes:
|
||
malware_cfg_list.append(file_name+".gexf")
|
||
if len(malware_cfg_list) == malware_num:
|
||
break
|
||
|
||
def load_data(paths, test_list=None,limit=None,normal=False):
|
||
print("正在加载二进制函数为向量：")
|
||
functions = []
|
||
with open(paths, 'r', encoding='utf-8') as f:
|
||
# 经下述操作后，reader成为了一个可以迭代行的文件
|
||
reader = csv.reader(f)
|
||
# 先拿出csv文件的首行（一般是基本名称说明的行），此时指针指向下一行
|
||
header = next(reader)
|
||
print(header)
|
||
for row in tqdm(reader):
|
||
func_name = row[0]
|
||
bytes = row[1]
|
||
fn = Function.load(bytes,normal)
|
||
# 在函数对象列表中添加fn
|
||
functions.append(fn)
|
||
# 如果测试函数列表存在，则判断加入的函数在不在测试列表中，不在的话退出
|
||
if test_list != None:
|
||
if len(test_list) == len(functions):
|
||
break
|
||
if func_name not in test_list:
|
||
continue
|
||
if limit and len(functions) >= limit:
|
||
break
|
||
# 在token列表中添加函数的
|
||
# tokens是每个函数中所有的操作符和操作数的列表
|
||
# 返回functions的列表与对应的token列表
|
||
return functions
|
||
|
||
|
||
#找到语境词与中心词、段向量
|
||
def preprocess(functions):
|
||
|
||
#上下文窗口大小
|
||
C = 1
|
||
#context_vec_list的每个成员由 段向量、前一条向量、后一条向量构成
|
||
context_vec_list=[]
|
||
center_vec_list = []
|
||
for i, fn in enumerate(functions):
|
||
hex2vec_list = fn.hex2vec_list
|
||
fun2vec = fn.fun2vec
|
||
j=C
|
||
while True:
|
||
center_word = hex2vec_list[j]
|
||
center_vec_list.append(center_word)
|
||
context_words = fun2vec + hex2vec_list[(j - C):j][0] + hex2vec_list[(j + 1):(j + C + 1)][0]
|
||
context_vec_list.append(context_words)
|
||
j+=1
|
||
if j >= len(hex2vec_list)-1:
|
||
break
|
||
|
||
return torch.tensor(context_vec_list), torch.tensor(center_vec_list)
|
||
|
||
def train(
|
||
functions,
|
||
model=None,
|
||
embedding_size=100,
|
||
batch_size=1024,
|
||
epochs=10,
|
||
calc_acc=False,
|
||
device='cuda:0',
|
||
mode='train',
|
||
callback=None,
|
||
learning_rate=0.02
|
||
):
|
||
vocab_size=len(functions[0].hex2vec_list[0])
|
||
# print(vocab_size)
|
||
# exit()
|
||
if mode == 'train':
|
||
if model is None:
|
||
model = ASM2VEC(vocab_size=vocab_size, function_size=len(functions), embedding_size=embedding_size).to(device)
|
||
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
|
||
elif mode == 'test':
|
||
if model is None:
|
||
raise ValueError("test mode required pretrained model")
|
||
optimizer = torch.optim.Adam(model.linear_f.parameters(), lr=learning_rate)
|
||
else:
|
||
raise ValueError("Unknown mode")
|
||
|
||
loader = DataLoader(AsmDataset(*preprocess(functions)), batch_size=batch_size, shuffle=True)
|
||
|
||
bar =tqdm(range(epochs),position=0)
|
||
for epoch in bar:
|
||
start = time.time()
|
||
loss_sum, loss_count, accs = 0.0, 0, []
|
||
model.train()
|
||
#inp是语境词向量，pos是中心词向量，neg是错误的中心词
|
||
# for i, (inp, pos) in enumerate(loader):
|
||
for i, (context_vec, center_vec) in enumerate(loader):
|
||
# print(context_vec)
|
||
# print(len(context_vec))
|
||
# print(center_vec)
|
||
# print(len(center_vec))
|
||
# exit()
|
||
loss = model(context_vec.to(device), center_vec.to(device))
|
||
loss_sum, loss_count = loss_sum + loss, loss_count + 1
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
|
||
if i == 0 and calc_acc:
|
||
probs = model.predict(context_vec.to(device), center_vec.to(device))
|
||
accs.append(accuracy(center_vec, probs))
|
||
|
||
if callback:
|
||
callback({
|
||
'model': model,
|
||
# 'tokens': tokens,
|
||
'epoch': epoch,
|
||
'time': time.time() - start,
|
||
'loss': loss_sum / loss_count,
|
||
'accuracy': torch.tensor(accs).mean() if calc_acc else None
|
||
})
|
||
# bar.set_postfix(loss=f"{loss_sum:.6f}", epoch=epoch + 1)
|
||
return model
|
||
|
||
def save_model(path, model, ):
|
||
torch.save({
|
||
'model_params': (
|
||
model.embeddings.num_embeddings,
|
||
model.embeddings_f.num_embeddings,
|
||
model.embeddings.embedding_dim
|
||
),
|
||
'model': model.state_dict(),
|
||
|
||
}, path)
|
||
|
||
def load_model(path, device='cpu'):
|
||
checkpoint = torch.load(path, map_location=device)
|
||
# tokens = Tokens()
|
||
# tokens.load_state_dict(checkpoint['tokens'])
|
||
model = ASM2VEC(*checkpoint['model_params'])
|
||
model.load_state_dict(checkpoint['model'])
|
||
model = model.to(device)
|
||
return model
|
||
|
||
def show_probs(x, y, probs, tokens, limit=None, pretty=False):
|
||
if pretty:
|
||
TL, TR, BL, BR = '┌', '┐', '└', '┘'
|
||
LM, RM, TM, BM = '├', '┤', '┬', '┴'
|
||
H, V = '─', '│'
|
||
arrow = ' ➔'
|
||
else:
|
||
TL = TR = BL = BR = '+'
|
||
LM = RM = TM = BM = '+'
|
||
H, V = '-', '|'
|
||
arrow = '->'
|
||
top = probs.topk(5)
|
||
for i, (xi, yi) in enumerate(zip(x, y)):
|
||
if limit and i >= limit:
|
||
break
|
||
xi, yi = xi.tolist(), yi.tolist()
|
||
print(TL + H * 42 + TR)
|
||
print(f'{V} {str(Instruction(tokens[xi[1]], tokens[xi[2:4]])):37} {V}')
|
||
print(f'{V} {arrow} {str(Instruction(tokens[yi[0]], tokens[yi[1:3]])):37} {V}')
|
||
print(f'{V} {str(Instruction(tokens[xi[4]], tokens[xi[5:7]])):37} {V}')
|
||
print(LM + H * 8 + TM + H * 33 + RM)
|
||
for value, index in zip(top.values[i], top.indices[i]):
|
||
if index in yi:
|
||
colorbegin, colorclear = '\033[92m', '\033[0m'
|
||
else:
|
||
colorbegin, colorclear = '', ''
|
||
print(f'{V} {colorbegin}{value*100:05.2f}%{colorclear} {V} {colorbegin}{tokens[index.item()].name:31}{colorclear} {V}')
|
||
print(BL + H * 8 + BM + H * 33 + BR)
|
||
|
||
def accuracy(y, probs):
|
||
y=y.type(torch.bool)
|
||
# print()
|
||
# print(probs)
|
||
# exit()
|
||
return torch.mean(torch.tensor([torch.sum(probs[i][yi]) for i, yi in enumerate(y)]))
|
||
|