diff --git a/asm2vec/utils.py b/asm2vec/utils.py index f854428..37a6ea6 100644 --- a/asm2vec/utils.py +++ b/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.""" diff --git a/debug_similarity_matrix.py b/debug_similarity_matrix.py new file mode 100644 index 0000000..eac28a2 --- /dev/null +++ b/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() diff --git a/generate_act_format.py b/generate_act_format.py index 417af31..4104159 100644 --- a/generate_act_format.py +++ b/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__": diff --git a/generate_deep_traces.py b/generate_deep_traces.py index eaaa1f3..6eb5a93 100644 --- a/generate_deep_traces.py +++ b/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) diff --git a/generate_similarity_matrix.py b/generate_similarity_matrix.py index 5774d99..e270095 100644 --- a/generate_similarity_matrix.py +++ b/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 - # (0, 2), then subtract 1 to get a score from range (-1, 1) + # Multiply the length factor and cosine similarity 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 + # 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"]: diff --git a/minor_patch_diff.py b/minor_patch_diff.py index 9c7d22c..ff48d1e 100644 --- a/minor_patch_diff.py +++ b/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], diff --git a/utils.py b/utils.py index 564b3cb..39484f5 100644 --- a/utils.py +++ b/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) diff --git a/vtable_alignment.py b/vtable_alignment.py index a190f02..fe9538b 100644 --- a/vtable_alignment.py +++ b/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 = []