338 lines
12 KiB
Python
338 lines
12 KiB
Python
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.
|
|
"""
|
|
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):
|
|
"""
|
|
Compares the OLD_TRACES and NEW_TRACES directories generated by the
|
|
`generate_deep_traces.py` script.
|
|
|
|
Creates a JSON OUTPUT_FILE containing best matches and opcodes with
|
|
ambiguous candidates that did not yield a definite match.
|
|
|
|
The format of the output is a list (all fields are optional):
|
|
|
|
\b
|
|
[
|
|
{
|
|
"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),
|
|
},
|
|
...
|
|
]
|
|
},
|
|
...
|
|
]
|
|
Example:
|
|
|
|
python traces_diff.py old-traces/ new-traces/ diff.json
|
|
"""
|
|
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()
|