diff --git a/asm2vec/utils.py b/asm2vec/utils.py index 2f77491..f854428 100644 --- a/asm2vec/utils.py +++ b/asm2vec/utils.py @@ -26,78 +26,71 @@ class TraceData: 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). + fn_idx: + Index of the Function. This is necessary to index into the + training model embeddings. Different switch cases may share + the same fn_idx because they may be textually identical but + have different constants. - idx: - Index of the ID/Function. This is necessary to index into the - training model embeddings. - - Opcode set: + Opcode Set/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: + Constants vector/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. + constants used in the trace. """ def __init__(self, tokens): self.tokens = tokens - self.traces = dict() + self.__traces = dict() """ - maps trace => ID. + Maps trace to trace idx + """ + + self.opcodes = dict() + """ + maps ptr_opcode => { fn_idx, constants_vector, [opcodes...] }. 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() + self.__opcode_sets = dict() """ maps ptr_opcode => [opcodes...] - See class docstriing for more details. """ + @property + def traces(self): + """ + A list of parsed traces in the TraceData. + + See class docstring for more details. + """ + return list(self.__traces.keys()) + 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) + if fn in self.__traces: + fn_idx = self.__traces[fn] else: - # Use ptr_opcode as an ID - id = ptr_opcode - - self.traces[fn] = id + fn_idx = len(self.__traces) + self.__traces[fn] = fn_idx 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) + self.opcodes[ptr_opcode] = { + "fn_idx": fn_idx, + "constants_vector": self.__get_constants_vector(text), + "opcodes": self.__opcode_sets[ptr_opcode], + } 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()} + self.__opcode_sets = {int(op): ops for op, ops in data.items()} @staticmethod def __read_trace_from_file(f): @@ -144,6 +137,9 @@ class TraceData: filenames += [Path(path)] trace_data = TraceData(tokens) + # Process opcode_sets.json first, save traces for later + + trace_files = dict() for filepath in filenames: filename = os.path.basename(filepath) file_split = os.path.splitext(filename) @@ -152,9 +148,12 @@ class TraceData: 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) + trace_files[ptr_opcode] = filepath + + for ptr_opcode, trace_file in trace_files.items(): + with open(trace_file) as f: + text = trace_data.__read_trace_from_file(f) + trace_data.__process_trace(ptr_opcode, text) return trace_data @@ -200,7 +199,7 @@ def train( learning_rate=0.02, ): """Trains the model on the provided trace data.""" - functions = trace_data.traces.keys() + functions = trace_data.traces tokens = trace_data.tokens if mode == "train": diff --git a/generate_similarity_matrix.py b/generate_similarity_matrix.py new file mode 100644 index 0000000..5774d99 --- /dev/null +++ b/generate_similarity_matrix.py @@ -0,0 +1,171 @@ +import click +import torch +from asm2vec.utils import ( + TraceData, + AsmDataset, + preprocess, + train, + save_model, + cosine_similarities, +) +from asm2vec.datatype import Tokens +import json + + +def length_heuristic(l0, l1, debug=False): + """ + Function length heuristic (since asm2vec is terrible at handling mismatched lengths) + + Returns a similarity metric in the range [0, 1] + """ + 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 full_similarity_matrix( + cosine_similarity_matrix, + old_trace_data: TraceData, + new_trace_data: TraceData, +): + """ + Generates a matrix comparing all old opcodes to all new opcodes. + + Returns (old opcodes, new_opcodes, similarity_matrix) + """ + old_fns = old_trace_data.traces + new_fns = new_trace_data.traces + + old_opcodes = [] + new_opcodes = [] + # Full similarity matrix mapping old_opcodes => new_opcodes + similarity_matrix = [] + + for old_data in old_trace_data.opcodes.values(): + similarities = [] + for new_data in new_trace_data.opcodes.values(): + old_idx = old_data["fn_idx"] + new_idx = new_data["fn_idx"] + # Use the length of the instructions for the length heuristic + l0 = len(old_fns[old_idx].insts) + l1 = len(new_fns[new_idx].insts) + length_factor = length_heuristic(l0, l1) + # Since cosine similarity is in the range (-1, 1), add 1 to push it + # into the range (0, 2). + cs = cosine_similarity_matrix[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 * cs - 1 + + # Now we copy this similarity value for all opcodes in the new + # switch case + for op in new_data["opcodes"]: + similarities.append(score) + # Now we copy this similarity mapping for all opcodes in the old + # switch case + for op in old_data["opcodes"]: + similarity_matrix.append(similarities) + + for old_data in old_trace_data.opcodes.values(): + for op in old_data["opcodes"]: + old_opcodes.append(op) + + for new_data in new_trace_data.opcodes.values(): + for op in new_data["opcodes"]: + new_opcodes.append(op) + + return (old_opcodes, new_opcodes, similarity_matrix) + + +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 generate_similarity_matrix(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 a pairwise similarity matrix of all + opcodes found. + + \b + { + "old_opcodes": (list of old opcodes indexing dimension 0), + "new_opcodes": (list of new opcodes indexing dimenision 1), + "matrix": (m by n array of floats: [[]]), + } + Example: + + python generate_similarity_matrix.py old-traces/ new-traces/ similarity.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("Computing full similarity matrix") + old_opcodes, new_opcodes, similarity_matrix = full_similarity_matrix( + csm, old_trace_data, new_trace_data + ) + + with open(output_file, "w+") as f: + json.dump( + { + "old_opcodes": old_opcodes, + "new_opcodes": new_opcodes, + "matrix": similarity_matrix, + }, + f, + indent=4, + ) + print_banner(f"Output written to {output_file}") + + +if __name__ == "__main__": + generate_similarity_matrix() diff --git a/traces_diff.py b/traces_diff.py deleted file mode 100644 index 9f6adbd..0000000 --- a/traces_diff.py +++ /dev/null @@ -1,337 +0,0 @@ -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()