From 9dc9416d4f875dfcc36ef30614e7add576732359 Mon Sep 17 00:00:00 2001
From: Flawed <33593723+ff14wed@users.noreply.github.com>
Date: Tue, 3 Oct 2023 00:59:44 -0700
Subject: [PATCH] Add ability to modify the similarity matrix to force matches

---
 debug_similarity_matrix.py    | 14 +++++++++++++-
 generate_similarity_matrix.py | 26 +++++++++++++++-----------
 vtable_alignment.py           | 32 +++++++++++++++++++++++++++++++-
 3 files changed, 59 insertions(+), 13 deletions(-)

diff --git a/debug_similarity_matrix.py b/debug_similarity_matrix.py
index 45b288b..0027767 100644
--- a/debug_similarity_matrix.py
+++ b/debug_similarity_matrix.py
@@ -15,7 +15,13 @@ from utils import HexIntParamType
     is_flag=True,
     help="Prints a column of the similarity matrix given a new opcode",
 )
-def debug_similarity_matrix(similarity_json_file, opcode, reverse):
+@click.option(
+    "--accept",
+    default=None,
+    help="Modifies the similarity matrix file by accepting the match between the opcode argument and the argument to this option",
+    type=HexIntParamType(),
+)
+def debug_similarity_matrix(similarity_json_file, opcode, reverse, accept):
     """
     Given an old opcode, debug prints a row of the similarity matrix generated
     from generate_similarity_matrix.py.
@@ -26,6 +32,12 @@ def debug_similarity_matrix(similarity_json_file, opcode, reverse):
 
     entries = dict()
     similarity = Similarity(similarity_json_file)
+
+    if accept is not None:
+        similarity.accept(opcode, accept)
+        similarity.write_to_file(similarity_json_file)
+        return
+
     if reverse:
         print("Checking column of matrix since --reverse was provided")
         for old_opcode in similarity.old_opcodes:
diff --git a/generate_similarity_matrix.py b/generate_similarity_matrix.py
index 1ccf86a..54b1b75 100644
--- a/generate_similarity_matrix.py
+++ b/generate_similarity_matrix.py
@@ -99,6 +99,20 @@ def print_banner(text):
     print("")
 
 
+def write_matrix_to_file(output_file, old_opcodes, new_opcodes, similarity_matrix):
+    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}")
+
+
 @click.command()
 @click.argument(
     "old_traces", type=click.Path(exists=True, file_okay=False, resolve_path=True)
@@ -166,17 +180,7 @@ def generate_similarity_matrix(old_traces, new_traces, output_file):
         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}")
+    write_matrix_to_file(output_file, old_opcodes, new_opcodes, similarity_matrix)
 
 
 if __name__ == "__main__":
diff --git a/vtable_alignment.py b/vtable_alignment.py
index 7eb2c64..2399bfd 100644
--- a/vtable_alignment.py
+++ b/vtable_alignment.py
@@ -4,6 +4,7 @@ import numpy as np
 
 from vtable_diff import extract_opcode_data
 from utils import eprint
+from generate_similarity_matrix import write_matrix_to_file
 
 
 def needleman_wunsch(old_seq, new_seq, similarity, gap_penalty):
@@ -103,6 +104,21 @@ class Similarity:
         j = self.new_opcodes[new_opcode]
         return self.matrix[i][j]
 
+    def accept(self, old_opcode, new_opcode):
+        if old_opcode not in self.old_opcodes:
+            self.warnings.add(
+                f"WARNING: Could not find old opcode {hex(old_opcode)} in similarity matrix"
+            )
+            return
+        if new_opcode not in self.new_opcodes:
+            self.warnings.add(
+                f"WARNING: Could not find new opcode {hex(new_opcode)} in similarity matrix"
+            )
+            return
+        i = self.old_opcodes[old_opcode]
+        j = self.new_opcodes[new_opcode]
+        self.matrix[i][j] = 1
+
     def get_confident_matches(self, threshold=0.1):
         """
         Returns matches in the form of [(old,new), ...] that are confidently
@@ -149,6 +165,14 @@ class Similarity:
         for warning in self.warnings:
             eprint(warning)
 
+    def write_to_file(self, output_file):
+        write_matrix_to_file(
+            output_file,
+            list(self.old_opcodes.keys()),
+            list(self.new_opcodes.keys()),
+            self.matrix.tolist(),
+        )
+
 
 @click.command()
 @click.argument(
@@ -220,7 +244,13 @@ def vtable_alignment(old_exe, new_exe, similarity_json_file):
     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)}")
+            mismatch_text = ""
+            if new is not None:
+                mismatch_text = hex(new)
+
+            eprint(
+                f"Mismatch detected! {hex(old)} => {hex(truth)}, got {mismatch_text}"
+            )
             mismatched_old.add(old)
             mismatched_new[truth] = Placeholder(old, truth)