diff --git a/debug_similarity_matrix.py b/debug_similarity_matrix.py index eac28a2..03e285a 100644 --- a/debug_similarity_matrix.py +++ b/debug_similarity_matrix.py @@ -29,7 +29,9 @@ def debug_similarity_matrix(similarity_json_file, opcode): print(f"Similarities for {hex(opcode)}") - for new_opcode, similarity in entries.items(): + entries = list(entries.items()) + entries.sort(key=lambda x: x[1], reverse=True) + for new_opcode, similarity in entries: print(f"\t{hex(new_opcode)} => {similarity}") if similarity > max_score: max_opcode = new_opcode diff --git a/generate_similarity_matrix.py b/generate_similarity_matrix.py index e270095..1ccf86a 100644 --- a/generate_similarity_matrix.py +++ b/generate_similarity_matrix.py @@ -9,8 +9,6 @@ from asm2vec.utils import ( from asm2vec.datatype import Tokens import json -import Levenshtein as lev - def length_heuristic(l0, l1, debug=False): """ diff --git a/vtable_alignment.py b/vtable_alignment.py index fe9538b..f3c4a50 100644 --- a/vtable_alignment.py +++ b/vtable_alignment.py @@ -1,5 +1,6 @@ import click import json +import numpy as np from vtable_diff import extract_opcode_data from utils import eprint @@ -64,6 +65,12 @@ def needleman_wunsch(old_seq, new_seq, similarity, gap_penalty): return alignment, mat[n][m] +class Placeholder: + def __init__(self, old, new): + self.old = old + self.new = new + + class Similarity: def __init__(self, similarity_json_file): with open(similarity_json_file) as f: @@ -75,10 +82,13 @@ class Similarity: self.new_opcodes = { opcode: idx for (idx, opcode) in enumerate(data["new_opcodes"]) } - self.matrix = data["matrix"] + self.matrix = np.array(data["matrix"]) self.warnings = set() def lookup(self, old_opcode, new_opcode): + if isinstance(old_opcode, Placeholder) or isinstance(new_opcode, Placeholder): + return -9999 + if old_opcode not in self.old_opcodes: self.warnings.add( f"WARNING: Could not find old opcode {hex(old_opcode)} in similarity matrix" @@ -93,6 +103,47 @@ class Similarity: j = self.new_opcodes[new_opcode] return self.matrix[i][j] + def get_confident_matches(self, threshold=0.1): + """ + Returns matches in the form of [(old,new), ...] that are confidently + above the score threshold and where all pairs in the matching prefer + each other over any other opcodes. + """ + scores = {} + new_opcodes = np.array([op for op in self.new_opcodes]) + for old_op, i in self.old_opcodes.items(): + top_n_idxs = np.argpartition(-self.matrix[i], 2)[:2] + scores[old_op] = [(new_opcodes[j], self.matrix[i][j]) for j in top_n_idxs] + scores[old_op].sort(key=lambda x: x[1], reverse=True) + + transposed_matrix = self.matrix.transpose() + rev_scores = {} + old_opcodes = np.array([op for op in self.old_opcodes]) + for new_op, j in self.new_opcodes.items(): + top_n_idxs = np.argpartition(-transposed_matrix[j], 2)[:2] + rev_scores[new_op] = [ + (old_opcodes[i], transposed_matrix[j][i]) for i in top_n_idxs + ] + rev_scores[new_op].sort(key=lambda x: x[1], reverse=True) + + # Accept matches if they pass the threshold in the old => new direction, + # and if the new => old direction is a best match + matches = [] + for old_op, top_matches in scores.items(): + new_op = top_matches[0][0] + if top_matches[0][1] - top_matches[1][1] >= threshold: + rev_top_matches = rev_scores[new_op] + if ( + rev_top_matches[0][0] == old_op + and rev_top_matches[0][1] - rev_top_matches[1][1] >= threshold + ): + matches.append((old_op, new_op)) + + return matches + + def clear_warnings(self): + self.warnings = set() + def print_warnings(self): for warning in self.warnings: eprint(warning) @@ -142,11 +193,61 @@ def vtable_alignment(old_exe, new_exe, similarity_json_file): new_seq = [opcode for opcode in new_opcodes_db.values()] similarity = Similarity(similarity_json_file) + + matches = similarity.get_confident_matches() + + # Ensure matches at least exist somewhere in the seq + old_seq_set = set(old_seq) + new_seq_set = set(new_seq) + matches = list( + filter(lambda x: x[0] in old_seq_set and x[1] in new_seq_set, matches) + ) + + eprint(f"Found {len(matches)} confident matches") + + eprint("Running initial alignment...") alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1) similarity.print_warnings() eprint(f"Alignment score: {score}") + # Check for mismatches + matched_old = {match[0]: match[1] for match in matches} + + mismatched_old = set() + mismatched_new = dict() + for old, new in alignment: + if old in matched_old and matched_old[old] != new: + truth = matched_old[old] + eprint(f"Mismatch detected! {hex(old)} => {hex(truth)}, got {hex(new)}") + mismatched_old.add(old) + mismatched_new[truth] = Placeholder(old, truth) + + eprint("Attempting to fix mismatches") + old_seq = list(filter(lambda x: x not in mismatched_old, old_seq)) + new_seq = list( + map(lambda x: mismatched_new[x] if x in mismatched_new else x, new_seq) + ) + + alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1) + + old_seq = [] + new_seq = [] + for old, target in alignment: + if isinstance(target, Placeholder): + old_seq.append(target.old) + new_seq.append(target.new) + else: + if old is not None: + old_seq.append(old) + if target is not None: + new_seq.append(target) + + similarity.clear_warnings() + alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1) + similarity.print_warnings() + eprint(f"Alignment score: {score}") + diff = [] for old, new in alignment: if old is None: