Add Asm2Vec method for comparing traces
This commit is contained in:
Flawed
@@ -0,0 +1,6 @@
|
||||
import importlib
|
||||
|
||||
__all__ = ["model", "datatype", "utils"]
|
||||
|
||||
for module in __all__:
|
||||
importlib.import_module(f".{module}", "asm2vec")
|
||||
@@ -0,0 +1,205 @@
|
||||
import torch
|
||||
import random
|
||||
import warnings
|
||||
import re
|
||||
|
||||
|
||||
class Token:
|
||||
def __init__(self, name, index):
|
||||
self.name = name
|
||||
self.index = index
|
||||
self.count = 1
|
||||
|
||||
def __str__(self):
|
||||
return self.name
|
||||
|
||||
|
||||
class Tokens:
|
||||
def __init__(self, name_to_index=None, tokens=None):
|
||||
self.name_to_index = name_to_index or {}
|
||||
self.tokens = tokens or []
|
||||
|
||||
def __getitem__(self, key):
|
||||
if type(key) is str:
|
||||
if self.name_to_index.get(key) is None:
|
||||
warnings.warn("Unknown token in training dataset")
|
||||
return self.tokens[self.name_to_index[""]]
|
||||
return self.tokens[self.name_to_index[key]]
|
||||
elif type(key) is int:
|
||||
return self.tokens[key]
|
||||
else:
|
||||
try:
|
||||
return [self[k] for k in key]
|
||||
except:
|
||||
raise ValueError
|
||||
|
||||
def load_state_dict(self, sd):
|
||||
self.name_to_index = sd["name_to_index"]
|
||||
self.tokens = sd["tokens"]
|
||||
|
||||
def state_dict(self):
|
||||
return {"name_to_index": self.name_to_index, "tokens": self.tokens}
|
||||
|
||||
def size(self):
|
||||
return len(self.tokens)
|
||||
|
||||
def add(self, names):
|
||||
if type(names) is not list:
|
||||
names = [names]
|
||||
for name in names:
|
||||
if name not in self.name_to_index:
|
||||
token = Token(name, len(self.tokens))
|
||||
self.name_to_index[name] = token.index
|
||||
self.tokens.append(token)
|
||||
else:
|
||||
self.tokens[self.name_to_index[name]].count += 1
|
||||
|
||||
def update(self, tokens_new):
|
||||
for token in tokens_new:
|
||||
if token.name not in self.name_to_index:
|
||||
token.index = len(self.tokens)
|
||||
self.name_to_index[token.name] = token.index
|
||||
self.tokens.append(token)
|
||||
else:
|
||||
self.tokens[self.name_to_index[token.name]].count += token.count
|
||||
|
||||
def precompute_weights(self, pos):
|
||||
"""
|
||||
This process actually takes a long time due to the size of the weights,
|
||||
so precompute them according to the shape of the dataset.
|
||||
"""
|
||||
from tqdm import tqdm
|
||||
|
||||
total = sum([token.count for token in self.tokens])
|
||||
token_weights = torch.zeros(len(self.tokens))
|
||||
for token in self.tokens:
|
||||
token_weights[token.index] = (token.count / total) ** 0.75
|
||||
weights = token_weights.repeat(pos.shape[0], 1)
|
||||
for i in tqdm(range(pos.shape[0])):
|
||||
for taken in pos[i]:
|
||||
weights[i][taken] = 0
|
||||
return weights
|
||||
|
||||
def sample(self, all_weights, batch_indices, num=5):
|
||||
weights = all_weights[batch_indices]
|
||||
return torch.multinomial(weights, num, replacement=False)
|
||||
|
||||
|
||||
class Function:
|
||||
def __init__(self, insts, blocks, meta):
|
||||
self.insts = insts
|
||||
self.blocks = blocks
|
||||
self.meta = meta
|
||||
|
||||
@classmethod
|
||||
def load(cls, text):
|
||||
"""
|
||||
gcc -S format compatiable
|
||||
"""
|
||||
label, labels, insts, blocks, meta = None, {}, [], [], {}
|
||||
for line in text.strip("\n").split("\n"):
|
||||
if line[0] in [" ", "\t"]:
|
||||
line = line.strip()
|
||||
# meta data
|
||||
if line[0] == ".":
|
||||
key, _, value = line[1:].strip().partition(" ")
|
||||
meta[key] = value
|
||||
# instruction
|
||||
else:
|
||||
inst = Instruction.load(line)
|
||||
insts.append(inst)
|
||||
if len(blocks) == 0 or blocks[-1].end():
|
||||
blocks.append(BasicBlock())
|
||||
# link prev and next block
|
||||
if len(blocks) > 1:
|
||||
blocks[-2].successors.add(blocks[-1])
|
||||
if label:
|
||||
labels[label], label = blocks[-1], None
|
||||
blocks[-1].add(inst)
|
||||
# label
|
||||
else:
|
||||
label = line.partition(":")[0]
|
||||
# link label
|
||||
for block in blocks:
|
||||
inst = block.insts[-1]
|
||||
if inst.is_jmp() and labels.get(inst.args[0]):
|
||||
block.successors.add(labels[inst.args[0]])
|
||||
# replace label with CONST
|
||||
for inst in insts:
|
||||
for i, arg in enumerate(inst.args):
|
||||
if labels.get(arg):
|
||||
inst.args[i] = "CONST"
|
||||
return cls(insts, blocks, meta)
|
||||
|
||||
def __hash__(self):
|
||||
return hash("\n".join((str(inst) for inst in self.insts)))
|
||||
|
||||
def __eq__(self, other):
|
||||
if isinstance(other, Function):
|
||||
a = "\n".join((str(inst) for inst in self.insts))
|
||||
b = "\n".join((str(inst) for inst in other.insts))
|
||||
return a == b
|
||||
return False
|
||||
|
||||
def tokens(self):
|
||||
return [token for inst in self.insts for token in inst.tokens()]
|
||||
|
||||
def random_walk(self, num=3):
|
||||
return [self._random_walk() for _ in range(num)]
|
||||
|
||||
def _random_walk(self):
|
||||
current, visited, seq = self.blocks[0], [], []
|
||||
while current not in visited:
|
||||
visited.append(current)
|
||||
seq += current.insts
|
||||
# no following block / hit return
|
||||
if len(current.successors) == 0 or current.insts[-1].op == "ret":
|
||||
break
|
||||
current = random.choice(list(current.successors))
|
||||
return seq
|
||||
|
||||
|
||||
class BasicBlock:
|
||||
def __init__(self):
|
||||
self.insts = []
|
||||
self.successors = set()
|
||||
|
||||
def add(self, inst):
|
||||
self.insts.append(inst)
|
||||
|
||||
def end(self):
|
||||
inst = self.insts[-1]
|
||||
return inst.is_jmp() or inst.op == "ret"
|
||||
|
||||
|
||||
class Instruction:
|
||||
def __init__(self, op, args):
|
||||
self.op = op
|
||||
self.args = args
|
||||
|
||||
def __str__(self):
|
||||
return f'{self.op} {", ".join([str(arg) for arg in self.args if str(arg)])}'
|
||||
|
||||
@classmethod
|
||||
def load(cls, text):
|
||||
text = text.strip().strip("bnd").strip() # get rid of BND prefix
|
||||
text = text.replace(" - ", " + ")
|
||||
text = re.sub(r"0x[0-9a-f]+", "CONST", text)
|
||||
text = re.sub(r"\*[0-9]", "*CONST", text)
|
||||
text = re.sub(r" [0-9]", " CONST", text)
|
||||
op, _, args = text.strip().partition(" ")
|
||||
if args:
|
||||
args = [arg.strip() for arg in args.split(",")]
|
||||
else:
|
||||
args = []
|
||||
args = (args + ["", ""])[:2]
|
||||
return cls(op, args)
|
||||
|
||||
def tokens(self):
|
||||
return [self.op] + self.args
|
||||
|
||||
def is_jmp(self):
|
||||
return "jmp" in self.op or self.op[0] == "j"
|
||||
|
||||
def is_call(self):
|
||||
return self.op == "call"
|
||||
@@ -0,0 +1,88 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
bce, sigmoid, softmax = nn.BCELoss(), nn.Sigmoid(), nn.Softmax(dim=1)
|
||||
|
||||
# TODO: Document how this shit works because wtf
|
||||
class ASM2VEC(nn.Module):
|
||||
def __init__(self, vocab_size, function_size, embedding_size):
|
||||
super(ASM2VEC, self).__init__()
|
||||
# Dictionary of the token embeddings: v_t
|
||||
self.embeddings = nn.Embedding(
|
||||
vocab_size,
|
||||
embedding_size,
|
||||
_weight=(torch.rand(vocab_size, embedding_size) - 0.5) / embedding_size / 2,
|
||||
)
|
||||
# Dictionary of the function embeddings, \theta_f_s
|
||||
self.embeddings_f = nn.Embedding(
|
||||
function_size,
|
||||
2 * embedding_size,
|
||||
_weight=(torch.rand(function_size, 2 * embedding_size) - 0.5)
|
||||
/ embedding_size
|
||||
/ 2,
|
||||
)
|
||||
# Dictionary of outputs: Transposed v'_t
|
||||
self.embeddings_r = nn.Embedding(
|
||||
vocab_size,
|
||||
2 * embedding_size,
|
||||
_weight=torch.zeros(vocab_size, 2 * embedding_size),
|
||||
)
|
||||
# Where the old embeddings are stored once the training step is done
|
||||
self.old_embeddings_f = None
|
||||
|
||||
def init_estimation_mode(self, function_size_new):
|
||||
device = self.embeddings.weight.device
|
||||
embedding_size = self.embeddings.embedding_dim
|
||||
|
||||
embedding_weights = self.embeddings.weight
|
||||
self.embeddings = nn.Embedding.from_pretrained(embedding_weights)
|
||||
|
||||
output_weights = self.embeddings_r.weight
|
||||
self.embeddings_r = nn.Embedding.from_pretrained(output_weights)
|
||||
|
||||
self.old_embeddings_f = self.embeddings_f
|
||||
self.embeddings_f = nn.Embedding(
|
||||
function_size_new,
|
||||
2 * embedding_size,
|
||||
_weight=(
|
||||
(torch.rand(function_size_new, 2 * embedding_size) - 0.5)
|
||||
/ embedding_size
|
||||
/ 2
|
||||
).to(device),
|
||||
)
|
||||
|
||||
def v(self, inp):
|
||||
# Retrieve the embeddings for all the context tokens
|
||||
e = self.embeddings(inp[:, 1:])
|
||||
# Retrieve the embedding for the function, \theta_f_s
|
||||
v_f = self.embeddings_f(inp[:, 0])
|
||||
# Calculate CT(in_(j-1))
|
||||
v_prev = torch.cat([e[:, 0], (e[:, 1] + e[:, 2]) / 2], dim=1)
|
||||
# Calculate CT(in_(j+1))
|
||||
v_next = torch.cat([e[:, 3], (e[:, 4] + e[:, 5]) / 2], dim=1)
|
||||
# delta(in_j, f_s) = 1/3 * (\theta_f_s + CT(in_(j-1)) + CT(in_(j+1)))
|
||||
v = ((v_f + v_prev + v_next) / 3).unsqueeze(2)
|
||||
return v
|
||||
|
||||
def forward(self, inp, pos, neg):
|
||||
device, batch_size = inp.device, inp.shape[0]
|
||||
v = self.v(inp)
|
||||
# negative sampling loss
|
||||
pred = torch.bmm(self.embeddings_r(torch.cat([pos, neg], dim=1)), v).squeeze()
|
||||
label = torch.cat(
|
||||
[torch.ones(batch_size, 1), torch.zeros(batch_size, neg.shape[1])], dim=1
|
||||
).to(device)
|
||||
return bce(sigmoid(pred), label)
|
||||
|
||||
def predict(self, inp, pos):
|
||||
device, batch_size = inp.device, inp.shape[0]
|
||||
v = self.v(inp)
|
||||
probs = torch.bmm(
|
||||
self.embeddings_r(
|
||||
torch.arange(self.embeddings_r.num_embeddings)
|
||||
.repeat(batch_size, 1)
|
||||
.to(device)
|
||||
),
|
||||
v,
|
||||
).squeeze(dim=2)
|
||||
return softmax(probs)
|
||||
@@ -0,0 +1,314 @@
|
||||
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
|
||||
|
||||
import re
|
||||
import json
|
||||
|
||||
CONSTANTS_RE = re.compile(r"(-? 0x[0-9a-f]+)|\*([0-9])| ([0-9])")
|
||||
|
||||
|
||||
class TraceData:
|
||||
"""
|
||||
A class that stores information about traces read from file.
|
||||
|
||||
Terminology:
|
||||
Trace/Function:
|
||||
The text produced by the generate_deep_traces.py script for each
|
||||
opcode switch case.
|
||||
|
||||
Token:
|
||||
A numerical representation of the each symbol that appears in the
|
||||
trace. The tokens dictionary should be shared across all traces to
|
||||
ensure the same representation matches.
|
||||
|
||||
ID:
|
||||
A representative identifier for the trace of a switch case. An ID
|
||||
may point to more than one switch case, if the traces of those
|
||||
switch cases are textually identical (ignoring constants).
|
||||
|
||||
idx:
|
||||
Index of the ID/Function. This is necessary to index into the
|
||||
training model embeddings.
|
||||
|
||||
Opcode set:
|
||||
A set of opcodes that a single switch case covers.
|
||||
|
||||
Pointer Opcode/ptr_opcode:
|
||||
A single opcode that represents the opcode set.
|
||||
|
||||
Constants vector:
|
||||
Since the training process removes constants from the input text,
|
||||
the constants vector is a "signature" generated by looking at
|
||||
constants used in the trace. These are correlated with ptr_opcodes
|
||||
and not IDs since constants can differ among cases represented by
|
||||
the same ID.
|
||||
"""
|
||||
|
||||
def __init__(self, tokens):
|
||||
self.tokens = tokens
|
||||
|
||||
self.traces = dict()
|
||||
"""
|
||||
maps trace => ID.
|
||||
See class docstring for more details.
|
||||
"""
|
||||
|
||||
self.ids = dict()
|
||||
"""
|
||||
maps ID => { idx, [ptr_opcodes...] }.
|
||||
See class docstring for more details.
|
||||
"""
|
||||
|
||||
self.constants_vectors = dict()
|
||||
"""
|
||||
maps ptr_opcode => constants_vector.
|
||||
See class docstriing for more details.
|
||||
"""
|
||||
|
||||
self.opcode_sets = dict()
|
||||
"""
|
||||
maps ptr_opcode => [opcodes...]
|
||||
See class docstriing for more details.
|
||||
"""
|
||||
|
||||
def __process_trace(self, ptr_opcode, text):
|
||||
fn = Function.load(text)
|
||||
if fn in self.traces:
|
||||
id = self.traces[fn]
|
||||
self.ids[id]["ptr_opcodes"].append(ptr_opcode)
|
||||
else:
|
||||
# Use ptr_opcode as an ID
|
||||
id = ptr_opcode
|
||||
|
||||
self.traces[fn] = id
|
||||
self.tokens.add(fn.tokens())
|
||||
self.ids[id] = {
|
||||
"idx": len(self.ids),
|
||||
"ptr_opcodes": [ptr_opcode],
|
||||
}
|
||||
|
||||
self.constants_vectors[ptr_opcode] = self.__get_constants_vector(text)
|
||||
|
||||
def __process_opcode_sets(self, opcode_sets_file):
|
||||
with open(opcode_sets_file) as f:
|
||||
data = json.load(f)
|
||||
self.opcode_sets = {int(op): ops for op, ops in data.items()}
|
||||
|
||||
@staticmethod
|
||||
def __read_trace_from_file(f):
|
||||
lines = f.readlines()
|
||||
normalized_lines = []
|
||||
for line in lines:
|
||||
normalized_lines.append(" " + line)
|
||||
# cap lines at 200 since Asm2Vec performance goes way down if
|
||||
# the text is too long
|
||||
return "".join(normalized_lines[:200])
|
||||
|
||||
@staticmethod
|
||||
def __get_constants_vector(trace):
|
||||
constants = []
|
||||
for line in trace.strip("\n").split("\n"):
|
||||
match = CONSTANTS_RE.search(line)
|
||||
if not match or not match.lastindex:
|
||||
continue
|
||||
|
||||
const_str = match.group(match.lastindex)
|
||||
const_str = "".join(const_str.split())
|
||||
const = int(const_str, 16)
|
||||
|
||||
# We only really care about constants less than 10000
|
||||
if abs(const) < 10000:
|
||||
constants.append(const)
|
||||
return constants
|
||||
|
||||
@staticmethod
|
||||
def load_data(paths, tokens):
|
||||
"""Reads traces from paths and returns TraceData for that path."""
|
||||
if type(paths) is not list:
|
||||
paths = [paths]
|
||||
|
||||
filenames = []
|
||||
for path in paths:
|
||||
if os.path.isdir(path):
|
||||
filenames += [
|
||||
Path(path) / filename
|
||||
for filename in sorted(os.listdir(path))
|
||||
if os.path.isfile(Path(path) / filename)
|
||||
]
|
||||
else:
|
||||
filenames += [Path(path)]
|
||||
|
||||
trace_data = TraceData(tokens)
|
||||
for filepath in filenames:
|
||||
filename = os.path.basename(filepath)
|
||||
file_split = os.path.splitext(filename)
|
||||
file_ext = file_split[-1]
|
||||
if file_ext == ".json" and filename == "opcode_sets.json":
|
||||
trace_data.__process_opcode_sets(filepath)
|
||||
elif file_ext == ".asm":
|
||||
ptr_opcode = int(file_split[0], base=16)
|
||||
with open(filepath) as f:
|
||||
text = trace_data.__read_trace_from_file(f)
|
||||
trace_data.__process_trace(ptr_opcode, text)
|
||||
|
||||
return trace_data
|
||||
|
||||
|
||||
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 index, self.x[index], self.y[index]
|
||||
|
||||
|
||||
def preprocess(functions, tokens):
|
||||
x, y = [], []
|
||||
for i, fn in enumerate(functions):
|
||||
for j in range(1, len(fn.insts) - 1):
|
||||
x.append(
|
||||
[i]
|
||||
+ [
|
||||
tokens[token].index
|
||||
for token in fn.insts[j - 1].tokens() + fn.insts[j + 1].tokens()
|
||||
]
|
||||
)
|
||||
y.append([tokens[token].index for token in fn.insts[j].tokens()])
|
||||
return torch.tensor(x), torch.tensor(y)
|
||||
|
||||
|
||||
def train(
|
||||
trace_data: TraceData,
|
||||
model=None,
|
||||
embedding_size=100,
|
||||
batch_size=1024,
|
||||
epochs=10,
|
||||
neg_sample_num=25,
|
||||
calc_acc=False,
|
||||
device="cpu",
|
||||
mode="train",
|
||||
callback=None,
|
||||
learning_rate=0.02,
|
||||
):
|
||||
"""Trains the model on the provided trace data."""
|
||||
functions = trace_data.traces.keys()
|
||||
tokens = trace_data.tokens
|
||||
|
||||
if mode == "train":
|
||||
if model is None:
|
||||
model = ASM2VEC(
|
||||
tokens.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 requires a pretrained model")
|
||||
optimizer = torch.optim.Adam(model.embeddings_f.parameters(), lr=learning_rate)
|
||||
else:
|
||||
raise ValueError("Unknown mode")
|
||||
|
||||
# Precompute the token weights so that they are cached for later use
|
||||
inp, pos = preprocess(functions, tokens)
|
||||
token_weights = tokens.precompute_weights(pos)
|
||||
loader = DataLoader(AsmDataset(inp, pos), batch_size=batch_size, shuffle=True)
|
||||
for epoch in range(epochs):
|
||||
start = time.time()
|
||||
loss_sum, loss_count, accs = 0.0, 0, []
|
||||
|
||||
model.train()
|
||||
"""
|
||||
Recall that the model consumes context tokens as input and uses
|
||||
embeddings for each output token as an output layer of the network. The
|
||||
model should output a high score for the token that correctly matches
|
||||
with the given context, and a low score for tokens that don't match the
|
||||
given context.
|
||||
|
||||
Our dataloader outputs samples in the form of
|
||||
batch_size x (batch_index, input_context, correct_token).
|
||||
"""
|
||||
for i, (batch_indices, inp, pos) in enumerate(loader):
|
||||
for j in range(pos.shape[1]):
|
||||
# Sample tokens that are not the positive token.
|
||||
neg = tokens.sample(token_weights, batch_indices, neg_sample_num)
|
||||
pos_token = torch.unsqueeze(pos[:, j], 1)
|
||||
loss = model(inp.to(device), pos_token.to(device), neg.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(inp.to(device), pos.to(device))
|
||||
accs.append(accuracy(pos, 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,
|
||||
}
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def save_model(path, model, tokens):
|
||||
torch.save(
|
||||
{
|
||||
"model_params": (
|
||||
model.embeddings.num_embeddings,
|
||||
model.embeddings_f.num_embeddings,
|
||||
model.embeddings.embedding_dim,
|
||||
),
|
||||
"model": model.state_dict(),
|
||||
"tokens": tokens.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, tokens
|
||||
|
||||
|
||||
def accuracy(y, probs):
|
||||
return torch.mean(torch.tensor([torch.sum(probs[i][yi]) for i, yi in enumerate(y)]))
|
||||
|
||||
|
||||
def cosine_similarities(model):
|
||||
"""
|
||||
Reads the old and new embeddings from the model and returns the pairwise
|
||||
cosine similarities between these embeddings.
|
||||
"""
|
||||
old_f = model.to("cpu").old_embeddings_f
|
||||
new_f = model.to("cpu").embeddings_f
|
||||
v_old = old_f(torch.tensor([i for i in range(old_f.num_embeddings)]))
|
||||
v_new = new_f(torch.tensor([i for i in range(new_f.num_embeddings)]))
|
||||
|
||||
cs_matrix = torch.nn.functional.cosine_similarity(
|
||||
v_old[:, :, None], v_new.t()[None, :, :]
|
||||
)
|
||||
|
||||
return cs_matrix.detach().numpy()
|
||||
+324
@@ -0,0 +1,324 @@
|
||||
import click
|
||||
import torch
|
||||
from asm2vec.utils import (
|
||||
TraceData,
|
||||
AsmDataset,
|
||||
preprocess,
|
||||
train,
|
||||
save_model,
|
||||
cosine_similarities,
|
||||
)
|
||||
from asm2vec.datatype import Tokens
|
||||
import json
|
||||
|
||||
import Levenshtein as lev
|
||||
|
||||
|
||||
class OpcodeMatcher:
|
||||
def __init__(
|
||||
self,
|
||||
cosine_similarity_matrix,
|
||||
old_trace_data: TraceData,
|
||||
new_trace_data: TraceData,
|
||||
):
|
||||
self.csm = cosine_similarity_matrix
|
||||
self.old_trace_data = old_trace_data
|
||||
self.new_trace_data = new_trace_data
|
||||
|
||||
self.old_functions = list(old_trace_data.traces.keys())
|
||||
self.new_functions = list(new_trace_data.traces.keys())
|
||||
|
||||
self.old_ptr_opcodes = self.__enumerate_ptr_opcodes(old_trace_data)
|
||||
self.new_ptr_opcodes = self.__enumerate_ptr_opcodes(new_trace_data)
|
||||
|
||||
self.cand_dict = None
|
||||
|
||||
def length_heuristic(self, old_idx, new_idx, debug=False):
|
||||
"""
|
||||
Function length heuristic (since asm2vec is terrible at handling mismatched lengths)
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
l0 = len(self.old_functions[old_idx].insts)
|
||||
l1 = len(self.new_functions[new_idx].insts)
|
||||
length_diff = abs(l0 - l1)
|
||||
# Weight mismatched lengths considerably lower, but clip factor to 0
|
||||
length_factor = max(1 - 4 * (length_diff / (l0 + l1)), 0)
|
||||
if debug:
|
||||
print("Length factor", l0, l1, length_factor)
|
||||
return length_factor
|
||||
|
||||
def constants_heuristic(self, old_opcode, new_opcode, debug=False):
|
||||
"""
|
||||
Constants vector heuristic.
|
||||
Runs a bit slow because it uses Levenshtein distance.
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
v0 = self.old_trace_data.constants_vectors[old_opcode]
|
||||
v1 = self.new_trace_data.constants_vectors[new_opcode]
|
||||
constants_diff = lev.distance(v0[:50], v1[:50])
|
||||
# Weight any constants differences harshly, but clip factor to 0
|
||||
constants_factor = max(1 - (constants_diff * 0.1), 0)
|
||||
if debug:
|
||||
print("Constants factor", constants_diff)
|
||||
print("v0", v0)
|
||||
print("v1", v1)
|
||||
return constants_factor
|
||||
|
||||
def case_length_heuristic(self, old_opcode, new_opcode, debug=False):
|
||||
"""
|
||||
Number of opcodes in case heuristic.
|
||||
|
||||
This one has questionable value because cases with different number of
|
||||
opcodes are already different enough from other cases.
|
||||
|
||||
Returns a similarity metric in the range [0, 1]
|
||||
"""
|
||||
n0 = len(self.old_trace_data.opcode_sets[old_opcode])
|
||||
n1 = len(self.new_trace_data.opcode_sets[new_opcode])
|
||||
case_length_diff = abs(n0 - n1)
|
||||
case_length_factor = max(1 - case_length_diff * 0.1, 0)
|
||||
if debug:
|
||||
print("Case length factor", n0, n1, case_length_diff)
|
||||
|
||||
return case_length_factor
|
||||
|
||||
@staticmethod
|
||||
def __enumerate_ptr_opcodes(trace_data: TraceData):
|
||||
ptr_opcodes = []
|
||||
for id, data in trace_data.ids.items():
|
||||
idx = data["idx"]
|
||||
for opcode in data["ptr_opcodes"]:
|
||||
ptr_opcodes.append((idx, opcode))
|
||||
return ptr_opcodes
|
||||
|
||||
def initialize_candidates(self):
|
||||
"""
|
||||
Generates a table of candidate matches between old pointer opcodes and
|
||||
new pointer opcodes.
|
||||
"""
|
||||
self.cand_dict = dict()
|
||||
|
||||
for (old_idx, old_opcode) in self.old_ptr_opcodes:
|
||||
candidates = []
|
||||
for (new_idx, new_opcode) in self.new_ptr_opcodes:
|
||||
length_factor = self.length_heuristic(old_idx, new_idx)
|
||||
constants_factor = self.constants_heuristic(old_opcode, new_opcode)
|
||||
case_length_factor = self.case_length_heuristic(old_opcode, new_opcode)
|
||||
# Since cosine similarity is in the range (-1, 1), add 1 to push it
|
||||
# into the range (0, 2).
|
||||
cs = self.csm[old_idx, new_idx] + 1
|
||||
|
||||
# Multiply all these factors together to yield some value in the range
|
||||
# (0, 2), then subtract 1 to get a score from range (-1, 1)
|
||||
score = length_factor * constants_factor * case_length_factor * cs - 1
|
||||
candidates.append((new_opcode, score))
|
||||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
if candidates[0][1] < -0.99:
|
||||
# If the top match is this low, then the heuristics screwed up the
|
||||
# candidates, so we'll have to go with just the cosine similarity metric
|
||||
candidates = [
|
||||
(new_opcode, self.csm[old_idx, new_idx])
|
||||
for (new_idx, new_opcode) in self.new_ptr_opcodes
|
||||
]
|
||||
candidates.sort(key=lambda x: x[1], reverse=True)
|
||||
candidates = candidates[:5]
|
||||
|
||||
self.cand_dict[old_opcode] = candidates
|
||||
|
||||
def accept_confident_matches(self, matches, threshold=0.1):
|
||||
"""
|
||||
Accepts matches for candidates where the score difference between the
|
||||
first and second best match is wider than the given threshold.
|
||||
"""
|
||||
num_new_matches = 0
|
||||
accepted_match_targets = set()
|
||||
unmatched = []
|
||||
|
||||
for opcode, candidates in self.cand_dict.items():
|
||||
if len(candidates) == 1 or (
|
||||
len(candidates) > 1
|
||||
and (candidates[0][1] - candidates[1][1] > threshold)
|
||||
):
|
||||
matches[opcode] = {
|
||||
"match": candidates[0][0],
|
||||
"score_lead": candidates[0][1] - candidates[1][1]
|
||||
if len(candidates) > 1
|
||||
else 0,
|
||||
}
|
||||
accepted_match_targets.add(candidates[0][0])
|
||||
num_new_matches += 1
|
||||
else:
|
||||
unmatched.append((opcode, candidates))
|
||||
|
||||
# Filter out match candidates that have already been matched
|
||||
new_candidates = dict()
|
||||
for opcode, candidates in unmatched:
|
||||
new_candidates[opcode] = [
|
||||
(cand_opcode, score)
|
||||
for (cand_opcode, score) in candidates
|
||||
if cand_opcode not in accepted_match_targets
|
||||
]
|
||||
|
||||
self.cand_dict = new_candidates
|
||||
return num_new_matches
|
||||
|
||||
def find_opcode_matches(self, threshold=0.1):
|
||||
"""
|
||||
Returns the best matches between old pointer opcodes and new pointer
|
||||
opcodes where the confidence is greater than the given threshold.
|
||||
"""
|
||||
matches = dict()
|
||||
self.initialize_candidates()
|
||||
num_new_matches = self.accept_confident_matches(matches, threshold)
|
||||
print("First pass added", num_new_matches, "matches")
|
||||
while num_new_matches > 0:
|
||||
num_new_matches = self.accept_confident_matches(matches, threshold)
|
||||
print("Added", num_new_matches, "additional matches")
|
||||
return matches
|
||||
|
||||
def find_matches_and_nonmatches(self):
|
||||
"""
|
||||
Returns the following information:
|
||||
|
||||
1. The best matches between old pointer opcodes and new pointer
|
||||
opcodes where the confidence is greater than the given threshold.
|
||||
2. Old opcodes for which a match could not be confidently found.
|
||||
3. New opcodes for which a match could not be confidently found.
|
||||
|
||||
The format of the output is a list (all fields are optional):
|
||||
[
|
||||
{
|
||||
"old": (list of opcodes in the switch case),
|
||||
"new": (list of opcodes in the switch case),
|
||||
"score_lead": (confidence above 2nd best match),
|
||||
"unknown": (true if a match was not made in this case),
|
||||
"candidates: [
|
||||
{
|
||||
"set": (list of opcodes in switch case),
|
||||
"score": (candidate score),
|
||||
}
|
||||
...
|
||||
]
|
||||
}
|
||||
]
|
||||
"""
|
||||
output = []
|
||||
|
||||
matches = self.find_opcode_matches()
|
||||
old_opcode_sets = self.old_trace_data.opcode_sets
|
||||
new_opcode_sets = self.new_trace_data.opcode_sets
|
||||
|
||||
for opcode, data in matches.items():
|
||||
output.append(
|
||||
{
|
||||
"old": [hex(old_opcode) for old_opcode in old_opcode_sets[opcode]],
|
||||
"new": [
|
||||
hex(new_opcode) for new_opcode in new_opcode_sets[data["match"]]
|
||||
],
|
||||
"score_lead": str(data["score_lead"]),
|
||||
}
|
||||
)
|
||||
|
||||
for opcode, candidates in self.cand_dict.items():
|
||||
output.append(
|
||||
{
|
||||
"old": [hex(old_opcode) for old_opcode in old_opcode_sets[opcode]],
|
||||
"candidates": [
|
||||
{
|
||||
"set": [
|
||||
hex(new_opcode)
|
||||
for new_opcode in new_opcode_sets[candidate]
|
||||
],
|
||||
"score": str(score),
|
||||
}
|
||||
for (candidate, score) in candidates
|
||||
if score > -1.0
|
||||
],
|
||||
"unknown": True,
|
||||
}
|
||||
)
|
||||
|
||||
unmatched_new_opcodes = set([opcode for (idx, opcode) in self.new_ptr_opcodes])
|
||||
for data in matches.values():
|
||||
if data["match"] in unmatched_new_opcodes:
|
||||
unmatched_new_opcodes.discard(data["match"])
|
||||
|
||||
for unmatched_opcode in unmatched_new_opcodes:
|
||||
output.append(
|
||||
{
|
||||
"new": [
|
||||
hex(new_opcode)
|
||||
for new_opcode in new_opcode_sets[unmatched_opcode]
|
||||
],
|
||||
"unknown": True,
|
||||
}
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def print_banner(text):
|
||||
print("")
|
||||
print(f"======= {text} =======")
|
||||
print("")
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.argument(
|
||||
"old_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
|
||||
)
|
||||
@click.argument(
|
||||
"new_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
|
||||
)
|
||||
@click.argument("output_file", type=click.Path(dir_okay=False, resolve_path=True))
|
||||
def traces_diff(old_traces, new_traces, output_file):
|
||||
tokens = Tokens()
|
||||
old_trace_data = TraceData.load_data(old_traces, tokens)
|
||||
new_trace_data = TraceData.load_data(new_traces, tokens)
|
||||
|
||||
opath = "model.pt"
|
||||
|
||||
def training_callback(context):
|
||||
progress = f'{context["epoch"]} | time = {context["time"]:.2f}, loss = {context["loss"]:.4f}'
|
||||
if context["accuracy"]:
|
||||
progress += f', accuracy = {context["accuracy"]:.4f}'
|
||||
print(progress)
|
||||
save_model(opath, context["model"], context["tokens"])
|
||||
|
||||
training_params = {
|
||||
"embedding_size": 100,
|
||||
"batch_size": 1024,
|
||||
"epochs": 20,
|
||||
"neg_sample_num": 25,
|
||||
"calc_acc": True,
|
||||
"device": "cuda" if torch.cuda.is_available() else "cpu",
|
||||
"callback": training_callback,
|
||||
"learning_rate": 0.02,
|
||||
}
|
||||
|
||||
print_banner("Training embeddings from scratch on old trace data")
|
||||
model = train(old_trace_data, **training_params)
|
||||
|
||||
# Prepare the model for new trace data and freeze all training from old trace data
|
||||
model.init_estimation_mode(len(new_trace_data.traces))
|
||||
|
||||
print_banner("Calculating embeddings for new trace data")
|
||||
model = train(new_trace_data, model=model, mode="test", **training_params)
|
||||
|
||||
print_banner("Calculating cosine similarities")
|
||||
csm = cosine_similarities(model)
|
||||
|
||||
print_banner("Calculating matches between opcodes")
|
||||
matcher = OpcodeMatcher(csm, old_trace_data, new_trace_data)
|
||||
compiled_data = matcher.find_matches_and_nonmatches()
|
||||
|
||||
with open(output_file, "w+") as f:
|
||||
json.dump(compiled_data, f, indent=2)
|
||||
print_banner(f"Output written to {output_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
traces_diff()
|
||||
Reference in New Issue
Block a user