94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
|
import torch
|
|||
|
|
import torch.nn as nn
|
|||
|
|
import click
|
|||
|
|
import asm2vec
|
|||
|
|
from utils2 import sigmoid, get_batches, compute_pca, get_dict
|
|||
|
|
from matplotlib import pyplot
|
|||
|
|
import numpy as np
|
|||
|
|
import os
|
|||
|
|
import random
|
|||
|
|
from asm2vec.get_opcode_vector import get_asm_input_vector,str_hex_to_bytes
|
|||
|
|
|
|||
|
|
def cosine_similarity(v1, v2):
|
|||
|
|
return (v1 @ v2 / (v1.norm() * v2.norm())).item()
|
|||
|
|
|
|||
|
|
def load_model(path="./asm2vec_checkpoints/model.pt"):
|
|||
|
|
|
|||
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|||
|
|
model = asm2vec.utils.load_model(path, device=device)
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
def func2vec1(model,hex_asm_list=["56a194382b56508b35cc912b56ffd6ff3584382b56ffd6a174382b5650ffd65ec3","56a194382b56508b35cc912b56ffd6ff3584382b56ffd6a174382b5650ffd65ec3"]):
|
|||
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|||
|
|
fun2vec_origin_list=[]
|
|||
|
|
for hex_asm in hex_asm_list:
|
|||
|
|
hex2vec_list = str_hex_to_bytes(hex_asm)
|
|||
|
|
hex2vec_list, opcode_oprand_seq = get_asm_input_vector(hex2vec_list)
|
|||
|
|
hex2vec_list=hex2vec_list
|
|||
|
|
fun2vec_origin = [0.0] * len(hex2vec_list[0])
|
|||
|
|
#开始对每一行的代码求平均值，得到函数的vec
|
|||
|
|
for i in hex2vec_list:
|
|||
|
|
for j in range(len(i)):
|
|||
|
|
fun2vec_origin[j] += i[j]
|
|||
|
|
opcode_seq_len=len(hex2vec_list)
|
|||
|
|
for i in range(len(fun2vec_origin)):
|
|||
|
|
fun2vec_origin[i] = fun2vec_origin[i] / opcode_seq_len
|
|||
|
|
fun2vec_origin=torch.tensor(fun2vec_origin).to(device)
|
|||
|
|
fun2vec_origin_list.append(fun2vec_origin)
|
|||
|
|
|
|||
|
|
fun2vec_origin_list = torch.tensor([item.cpu().detach().numpy() for item in fun2vec_origin_list]).cuda()
|
|||
|
|
|
|||
|
|
# print(fun2vec_origin_list)
|
|||
|
|
# exit()
|
|||
|
|
embedding_func_vec = model.to(device).linear_f(torch.tensor(fun2vec_origin_list).to(device)).clone().detach().requires_grad_(True)
|
|||
|
|
# print(embedding_func_vec)
|
|||
|
|
return embedding_func_vec
|
|||
|
|
|
|||
|
|
def func2vec(model,hex_asm="56a194382b56508b35cc912b56ffd6ff3584382b56ffd6a174382b5650ffd65ec3"):
|
|||
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|||
|
|
# device="cpu"
|
|||
|
|
# print(device)
|
|||
|
|
# device="cpu"
|
|||
|
|
# exit()
|
|||
|
|
hex2vec_list = str_hex_to_bytes(hex_asm)
|
|||
|
|
|
|||
|
|
|
|||
|
|
hex2vec_list, opcode_oprand_seq = get_asm_input_vector(hex2vec_list)
|
|||
|
|
hex2vec_list=hex2vec_list
|
|||
|
|
# print(hex2vec_list)
|
|||
|
|
# exit()
|
|||
|
|
# print(hex2vec_list)
|
|||
|
|
# print(hex2vec_list)
|
|||
|
|
# exit()
|
|||
|
|
fun2vec_origin = [0.0] * len(hex2vec_list[0])
|
|||
|
|
|
|||
|
|
# print(hex2vec_list)
|
|||
|
|
# print(fun2vec_origin)
|
|||
|
|
# print(torch.tensor([0,0,0])+torch.tensor([12,2,3]))
|
|||
|
|
# exit()
|
|||
|
|
#开始对每一行的代码求平均值，得到函数的vec
|
|||
|
|
for i in hex2vec_list:
|
|||
|
|
for j in range(len(i)):
|
|||
|
|
fun2vec_origin[j] += i[j]
|
|||
|
|
# print(fun2vec_origin)
|
|||
|
|
# exit()
|
|||
|
|
opcode_seq_len=len(hex2vec_list)
|
|||
|
|
for i in range(len(fun2vec_origin)):
|
|||
|
|
fun2vec_origin[i] = fun2vec_origin[i] / opcode_seq_len
|
|||
|
|
# print(fun2vec_origin)
|
|||
|
|
fun2vec_origin=torch.tensor(fun2vec_origin).to(device)
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
embedding_func_vec = model.to(device).linear_f(fun2vec_origin)
|
|||
|
|
# print(embedding_func_vec)
|
|||
|
|
# exit()
|
|||
|
|
return embedding_func_vec
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
model= load_model(path="./asm2vec_checkpoints/model_100.pt")
|
|||
|
|
|
|||
|
|
# model=
|
|||
|
|
# func2vec(model,hex_asm="f044014910488b81e00000004885c07404")
|
|||
|
|
func2vec(model)
|
|||
|
|
# func2vec(model,hex_asm="56a194382b56508b35cc912b56ffd6ff3584382b56ffd6a174382b5650ffd65ec3")
|