diff --git a/asm2vec/__init__.py b/asm2vec/__init__.py new file mode 100644 index 0000000..0cfb3f3 --- /dev/null +++ b/asm2vec/__init__.py @@ -0,0 +1,6 @@ +import importlib + +__all__ = ["model", "datatype", "utils"] + +for module in __all__: + importlib.import_module(f".{module}", "asm2vec") diff --git a/asm2vec/datatype.py b/asm2vec/datatype.py new file mode 100644 index 0000000..e36dc98 --- /dev/null +++ b/asm2vec/datatype.py @@ -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" diff --git a/asm2vec/model.py b/asm2vec/model.py new file mode 100644 index 0000000..18923bc --- /dev/null +++ b/asm2vec/model.py @@ -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) diff --git a/asm2vec/utils.py b/asm2vec/utils.py new file mode 100644 index 0000000..2f77491 --- /dev/null +++ b/asm2vec/utils.py @@ -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() diff --git a/traces_diff.py b/traces_diff.py new file mode 100644 index 0000000..7f71624 --- /dev/null +++ b/traces_diff.py @@ -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()