diff --git a/sanity_check.py b/sanity_check.py index 81e8ee5..50af5ce 100644 --- a/sanity_check.py +++ b/sanity_check.py @@ -31,15 +31,23 @@ def sanity_check(vtable_diff, minor_patch_diff): diff1 = load_diff_file(vtable_diff) diff2 = load_diff_file(minor_patch_diff) - for (old_opcode, new_opcodes) in diff2.items(): + all_good = True + for old_opcode, new_opcodes in diff2.items(): if old_opcode not in diff1: if len(new_opcodes) < 50: print(f"Missing old opcode in vtable diff: {old_opcode}") + all_good = False continue other_new_opcodes = diff1[old_opcode] if list(other_new_opcodes)[0] not in new_opcodes: print(f"vtable diff mismatch for case {old_opcode} => {other_new_opcodes}") + all_good = False + + if all_good: + print("Sanity check OK") + else: + print("Errors detected") if __name__ == "__main__": diff --git a/vtable_alignment.py b/vtable_alignment.py new file mode 100644 index 0000000..a190f02 --- /dev/null +++ b/vtable_alignment.py @@ -0,0 +1,157 @@ +import click +import json + +from vtable_diff import extract_opcode_data +from utils import eprint + + +def needleman_wunsch(old_seq, new_seq, similarity, gap_penalty): + """ + The Needleman-Wunsch algorithm adapted from + https://github.com/farhanma/pyseq/blob/master/functions1_3.py + + Returns (alignment, alignment_score) + """ + + # Stage 1: Create a zero matrix and fills it via algorithm + n, m = len(old_seq), len(new_seq) + mat = [] + for i in range(n + 1): + mat.append([0] * (m + 1)) + for j in range(m + 1): + mat[0][j] = gap_penalty * j + for i in range(n + 1): + mat[i][0] = gap_penalty * i + for i in range(1, n + 1): + for j in range(1, m + 1): + 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, + mat[i - 1][j] + gap_penalty, + ) + + # Stage 2: Computes the final alignment, by backtracking through matrix + alignment = [] + i, j = n, m + while i and j: + score, scoreDiag, scoreUp, scoreLeft = ( + mat[i][j], + mat[i - 1][j - 1], + mat[i - 1][j], + mat[i][j - 1], + ) + if score == scoreDiag + similarity.lookup(old_seq[i - 1], new_seq[j - 1]): + alignment.append((old_seq[i - 1], new_seq[j - 1])) + i -= 1 + j -= 1 + elif score == scoreUp + gap_penalty: + alignment.append((old_seq[i - 1], None)) + i -= 1 + elif score == scoreLeft + gap_penalty: + alignment.append((None, new_seq[j - 1])) + j -= 1 + while i: + alignment.append((old_seq[i - 1], None)) + i -= 1 + while j: + alignment.append((None, new_seq[j - 1])) + j -= 1 + + # Since we were backtracking, we reverse the collected alignment + alignment.reverse() + + return alignment, mat[n][m] + + +class Similarity: + def __init__(self, similarity_json_file): + with open(similarity_json_file) as f: + data = json.load(f) + + self.old_opcodes = { + opcode: idx for (idx, opcode) in enumerate(data["old_opcodes"]) + } + self.new_opcodes = { + opcode: idx for (idx, opcode) in enumerate(data["new_opcodes"]) + } + self.__matrix = data["matrix"] + + def lookup(self, old_opcode, new_opcode): + if old_opcode not in self.old_opcodes: + eprint( + f"WARNING: Could not find old opcode {hex(old_opcode)} in similarity matrix" + ) + return 0 + if new_opcode not in self.new_opcodes: + eprint( + 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] + + +@click.command() +@click.argument( + "old_exe", type=click.Path(exists=True, dir_okay=False, resolve_path=True) +) +@click.argument( + "new_exe", type=click.Path(exists=True, dir_okay=False, resolve_path=True) +) +@click.argument( + "similarity_json_file", + type=click.Path(exists=True, dir_okay=False, resolve_path=True), +) +def vtable_alignment(old_exe, new_exe, similarity_json_file): + """ + A more generalized version of vtable_diff. Generates an opcode + diff file by running a sequence alignment algorithm and attempting + to find the optimal global alignment of the vtable opcodes + from different exe versions. + + Requires a similarity matrix generated from generate_similarity_matrix.py. + + This script outputs to stdout, so pipe it to a json file. + + The format of the output is a list (all fields are optional): + + \b + [ + { + "old": [opcode], + "new": [opcode], + }, + ... + ] + + Example: + + python vtable_alignment.py ffxiv_dx11.old.exe ffxiv_dx11.new.exe similarity.json > diff.json + """ + old_opcodes_db = extract_opcode_data(old_exe) + new_opcodes_db = extract_opcode_data(new_exe) + + old_seq = [opcode for opcode in old_opcodes_db.values()] + new_seq = [opcode for opcode in new_opcodes_db.values()] + + similarity = Similarity(similarity_json_file) + alignment, score = needleman_wunsch(old_seq, new_seq, similarity, -1) + eprint(f"Alignment score: {score}") + + diff = [] + for old, new in alignment: + if old is None: + eprint("New opcode did not find matching old one:", hex(new)) + diff.append({"old": [], "new": [hex(new)]}) + elif new is None: + eprint("Old opcode did not find matching new one:", hex(old)) + diff.append({"old": [hex(old)], "new": []}) + else: + diff.append({"old": [hex(old)], "new": [hex(new)]}) + + print(json.dumps(diff, indent=2)) + + +if __name__ == "__main__": + vtable_alignment()