Bunch of miscellaneous updates

I've dropped the ball on making smaller commits so now here's a big one

- Update sig for post-6.40 exes
- Some packet handlers have switches for packet sizes. Added a heuristic
  to capture this information
- Consolidated vtable alignment warnings
- Added script to debug similarity matrix

asm2vec/utils.py

@@ -10,6 +10,7 @@ import re
 import json
 
 CONSTANTS_RE = re.compile(r"(-? 0x[0-9a-f]+)|\*([0-9])| ([0-9])")
+PACKET_SIZE_HINT_RE = re.compile(r"mov qword \[rsp \+ 0x20\], (0x[0-9a-f]+)")
 
 
 class TraceData:
@@ -84,6 +85,7 @@ class TraceData:
         self.opcodes[ptr_opcode] = {
             "fn_idx": fn_idx,
             "constants_vector": self.__get_constants_vector(text),
+            "packet_size_hint": self.__get_packet_size_hint(text),
             "opcodes": self.__opcode_sets[ptr_opcode],
         }
 
@@ -119,6 +121,37 @@ class TraceData:
                 constants.append(const)
         return constants
 
+    @staticmethod
+    def __get_packet_size_hint(trace):
+        """
+        Dumb way of getting the packet size hint from whatever handlers
+        call them
+        """
+        call0_lines = []
+        call0_found = False
+
+        for line in trace.strip("\n").split("\n"):
+            if not call0_found:
+                if "CALL0" in line:
+                    call0_found = True
+                else:
+                    continue
+            if "CALL0_END" in line:
+                break
+            else:
+                call0_lines.append(line)
+
+        for line in call0_lines:
+            match = PACKET_SIZE_HINT_RE.search(line)
+            if match and match.lastindex:
+                const_str = match.group(match.lastindex)
+                const_str = "".join(const_str.split())
+                const = int(const_str, 16)
+                return const
+
+        # Couldn't find anything, so just return 0
+        return 0
+
     @staticmethod
     def load_data(paths, tokens):
         """Reads traces from paths and returns TraceData for that path."""

debug_similarity_matrix.py

@@ -0,0 +1,43 @@
+import click
+from vtable_alignment import Similarity
+
+from utils import HexIntParamType
+
+
+@click.command()
+@click.argument(
+    "similarity_json_file",
+    type=click.Path(exists=True, dir_okay=False, resolve_path=True),
+)
+@click.argument("opcode", type=HexIntParamType())
+def debug_similarity_matrix(similarity_json_file, opcode):
+    """
+    Given an old opcode, debug prints a row of the similarity matrix generated
+    from generate_similarity_matrix.py.
+
+    Example:
+    python debug_similarity_matrix.py similarity.json 0x200
+    """
+
+    entries = dict()
+    similarity = Similarity(similarity_json_file)
+    for new_opcode in similarity.new_opcodes:
+        entries[new_opcode] = similarity.lookup(opcode, new_opcode)
+
+    max_opcode = 0
+    max_score = -9999
+
+    print(f"Similarities for {hex(opcode)}")
+
+    for new_opcode, similarity in entries.items():
+        print(f"\t{hex(new_opcode)} => {similarity}")
+        if similarity > max_score:
+            max_opcode = new_opcode
+            max_score = similarity
+
+    print("Best match")
+    print(f"{hex(max_opcode)} => {max_score}")
+
+
+if __name__ == "__main__":
+    debug_similarity_matrix()

generate_act_format.py

@@ -53,7 +53,9 @@ def generate_act_format(opcodes_file):
 
         opcode_name = match_groups[0][0].strip()
         opcode_val = match_groups[0][1]
-        if " or " in opcode_val:
+        if opcode_val == "UNKNOWN":
+            opcodes = []
+        elif " or " in opcode_val:
             opcodes = [int(v, 16) for v in opcode_val.split(" or ")]
         else:
             opcodes = [int(opcode_val, 16)]
@@ -67,6 +69,8 @@ def generate_act_format(opcodes_file):
             print(f"{desired}|{opcodes[0]:x}")
         elif len(opcodes) > 1:
             print(f'{desired}|{[f"{opcode:x}" for opcode in opcodes]}')
+        else:
+            print(f"{desired}|???")
 
 
 if __name__ == "__main__":

generate_deep_traces.py

@@ -2,7 +2,11 @@ import click
 import json
 import pathlib
 
-from minor_patch_diff import get_opcode_offset, get_longest_switch, ZONE_PROTO_DOWN_SIG
+from minor_patch_diff import (
+    get_opcode_offset,
+    get_longest_switch,
+    get_zone_proto_down_sig,
+)
 
 
 class RefNode:
@@ -138,7 +142,7 @@ def extract_opcode_data(exe_file):
 
     sync_r2_output(r2)
 
-    p = create_r2_byte_pattern(ZONE_PROTO_DOWN_SIG)
+    p = create_r2_byte_pattern(get_zone_proto_down_sig(exe_file))
     target = r2.cmd(f"/x {p}").split()[0]  # Find byte pattern
     packet_handler_ea = int(target, 16)
 

generate_similarity_matrix.py

@@ -2,8 +2,6 @@ import click
 import torch
 from asm2vec.utils import (
     TraceData,
-    AsmDataset,
-    preprocess,
     train,
     save_model,
     cosine_similarities,
@@ -11,6 +9,8 @@ from asm2vec.utils import (
 from asm2vec.datatype import Tokens
 import json
 
+import Levenshtein as lev
+
 
 def length_heuristic(l0, l1, debug=False):
     """
@@ -53,14 +53,28 @@ def full_similarity_matrix(
             l0 = len(old_fns[old_idx].insts)
             l1 = len(new_fns[new_idx].insts)
             length_factor = length_heuristic(l0, l1)
+
+            h0 = old_data["packet_size_hint"]
+            h1 = new_data["packet_size_hint"]
+
+            packet_size_factor = 0
+            if h0 == h1 and h0 != 0:
+                packet_size_factor = 0.5
+            else:
+                packet_size_factor = -0.5
             # 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
+            # Multiply the length factor and cosine similarity together to
-            # (0, 2), then subtract 1 to get a score from range (-1, 1)
+            # yield some value in the range (0, 2), then subtract 1 to get a
+            # score from range (-1, 1)
             score = length_factor * cs - 1
 
+            # Add or subtract score depending on the packet size matching
+            # Also clamp value to between (-1, 1)
+            score = max(min(score + packet_size_factor, 1.0), -1.0)
+
             # Now we copy this similarity value for all opcodes in the new
             # switch case
             for op in new_data["opcodes"]:

minor_patch_diff.py

@@ -2,12 +2,19 @@ import click
 import json
 import re
 
-ZONE_PROTO_DOWN_SIG = "48 89 ? 24 ? ? 48 83 EC 50 8B F2 49 8B"
-
 fucked_distance = 0xFFFFFFFF
 max_size_diff = 10
 
 
+def get_zone_proto_down_sig(exe_file: str):
+    res = re.match(".*ffxiv_dx11\.(.*)\.exe", exe_file)
+    ver = res.group(1)
+    if ver == "6.40":
+        return "40 53 56 48 81 EC ? ? ? ? 48 8B 05 ? ? ? ? 48 33 C4 48 89 44 24 ? 8B F2"
+    else:
+        return "48 89 ? 24 ? ? 48 83 EC 50 8B F2 49 8B"
+
+
 def get_opcode_offset(r2):
     orig_loc = r2.cmd("s")  # Save original spot
     r2.cmd("aei")  # Initialize ESIL VM
@@ -15,6 +22,7 @@ def get_opcode_offset(r2):
     r2.cmd("aeip")  # Initialize ESIL VM IP to curseek
 
     r2.cmd("aecc")  # continue until call
+    r2.cmd("aer rax=0x0")  # set rax to 0
     r2.cmd('"aesue rax,0x0,>"')  # continue until rax changes?
     r2.cmd("aer rdx=0x200")  # set rdx to some arbitrary number
     r2.cmd("aeso")  # step
@@ -89,7 +97,7 @@ def extract_opcode_data(exe_file):
 
     sync_r2_output(r2)
 
-    p = create_r2_byte_pattern(ZONE_PROTO_DOWN_SIG)
+    p = create_r2_byte_pattern(get_zone_proto_down_sig(exe_file))
     target = r2.cmd(f"/x {p}").split()[0]  # Find byte pattern
     packet_handler_ea = int(target, 16)
 
@@ -216,7 +224,7 @@ def minor_patch_diff(old_exe, new_exe):
     opcodes_found = find_opcode_matches(old_opcodes_db, new_opcodes_db)
     opcodes_object = []
 
-    for (old, new) in opcodes_found:
+    for old, new in opcodes_found:
         opcodes_object.append(
             {
                 "old": [hex(o) for o in old],

utils.py

@@ -1,6 +1,7 @@
 import sys
 import time
 import os
+import click
 
 
 def eprint(*args, **kwargs):
@@ -38,3 +39,18 @@ def sync_r2_output(r2):
     output = r2.cmd(f"?vi 123").strip()
     if output != "123":
         raise Exception("R2 state never got synced")
+
+
+class HexIntParamType(click.ParamType):
+    name = "integer"
+
+    def convert(self, value, param, ctx):
+        if isinstance(value, int):
+            return value
+
+        try:
+            if value[:2].lower() == "0x":
+                return int(value[2:], 16)
+            return int(value, 16)
+        except ValueError:
+            self.fail(f"{value!r} is not a valid hex integer", param, ctx)

vtable_alignment.py

@@ -24,6 +24,7 @@ def needleman_wunsch(old_seq, new_seq, similarity, gap_penalty):
         mat[i][0] = gap_penalty * i
     for i in range(1, n + 1):
         for j in range(1, m + 1):
+            # max(Match, Insertion, Deletion)
             mat[i][j] = max(
                 mat[i - 1][j - 1] + similarity.lookup(old_seq[i - 1], new_seq[j - 1]),
                 mat[i][j - 1] + gap_penalty,
@@ -74,22 +75,27 @@ class Similarity:
         self.new_opcodes = {
             opcode: idx for (idx, opcode) in enumerate(data["new_opcodes"])
         }
-        self.__matrix = data["matrix"]
+        self.matrix = data["matrix"]
+        self.warnings = set()
 
     def lookup(self, old_opcode, new_opcode):
         if old_opcode not in self.old_opcodes:
-            eprint(
+            self.warnings.add(
                 f"WARNING: Could not find old opcode {hex(old_opcode)} in similarity matrix"
             )
             return 0
         if new_opcode not in self.new_opcodes:
-            eprint(
+            self.warnings.add(
                 f"WARNING: Could not find new opcode {hex(new_opcode)} in similarity matrix"
             )
             return 0
         i = self.old_opcodes[old_opcode]
         j = self.new_opcodes[new_opcode]
-        return self.__matrix[i][j]
+        return self.matrix[i][j]
+
+    def print_warnings(self):
+        for warning in self.warnings:
+            eprint(warning)
 
 
 @click.command()
@@ -137,6 +143,8 @@ def vtable_alignment(old_exe, new_exe, similarity_json_file):
 
     similarity = Similarity(similarity_json_file)
     alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1)
+
+    similarity.print_warnings()
     eprint(f"Alignment score: {score}")
 
     diff = []