Rename traces_diff -> generate_similarity_matrix

Rewrite the diff generation to generate a similarity
matrix instead

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()