From 4039ca24a05dfa9c3fd6ceeee9b14604fae78838 Mon Sep 17 00:00:00 2001
From: Flawed <33593723+ff14wed@users.noreply.github.com>
Date: Tue, 12 Nov 2024 00:42:56 -0800
Subject: [PATCH] Update switch finding code to find the correct one

---
 generate_deep_traces.py |  5 +++--
 minor_patch_diff.py     | 13 ++++++++++---
 vtable_diff.py          |  7 +++++--
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/generate_deep_traces.py b/generate_deep_traces.py
index 6eb5a93..59f7244 100644
--- a/generate_deep_traces.py
+++ b/generate_deep_traces.py
@@ -4,7 +4,7 @@ import pathlib
 
 from minor_patch_diff import (
     get_opcode_offset,
-    get_longest_switch,
+    get_correct_switch,
     get_zone_proto_down_sig,
 )
 
@@ -163,7 +163,8 @@ def extract_opcode_data(exe_file):
     fn_graph = r2.cmdj(f"pdrj")
 
     ## STEP 4: Process data
-    packet_handler_switch = get_longest_switch(switch_cases)
+    switch_ea, packet_handler_switch = get_correct_switch(packet_handler_ea, switch_cases)
+    eprint(f"  Found switch at {switch_ea}")
 
     opcodes_db, blocks = generate_opcodes_db(
         r2, packet_handler_switch, opcode_offset, fn_graph
diff --git a/minor_patch_diff.py b/minor_patch_diff.py
index 1cd18cf..e8d6251 100644
--- a/minor_patch_diff.py
+++ b/minor_patch_diff.py
@@ -42,7 +42,7 @@ def get_opcode_offset(r2):
     return opcode_offset
 
 
-def get_longest_switch(switch_cases):
+def get_correct_switch(approx_ea, switch_cases):
     switches = dict()
 
     pattern = re.compile("case\.(0x[0-9a-fA-F]+)\.(\d+)")
@@ -61,12 +61,17 @@ def get_longest_switch(switch_cases):
                 }
             switches[switch_ea][case_ea]["opcodes"].append(match[2])
 
+    found_switch = None
     longest_switch = dict()
     for switch_ea in switches:
+        int_switch_ea = int(switch_ea, 16)
+        if int_switch_ea < approx_ea or int_switch_ea > approx_ea+0x100:
+            continue
         if len(switches[switch_ea].keys()) > len(longest_switch):
+            found_switch = switch_ea
             longest_switch = switches[switch_ea]
 
-    return longest_switch
+    return found_switch, longest_switch
 
 
 def get_block_sizes(blocks):
@@ -126,7 +131,9 @@ def extract_opcode_data(exe_file):
     eprint(f"  Grabbed blocks from packet handler")
 
     ## STEP 4: Process data
-    packet_handler_switch = get_longest_switch(switch_cases)
+    switch_ea, packet_handler_switch = get_correct_switch(packet_handler_ea, switch_cases)
+    eprint(f"  Found switch at {switch_ea}")
+
     block_sizes = get_block_sizes(blocks)
     opcodes_db = generate_opcodes_db(
         packet_handler_ea, packet_handler_switch, opcode_offset, block_sizes
diff --git a/vtable_diff.py b/vtable_diff.py
index 54664b4..91c9572 100644
--- a/vtable_diff.py
+++ b/vtable_diff.py
@@ -1,7 +1,7 @@
 import click
 import json
 
-from minor_patch_diff import get_longest_switch
+from minor_patch_diff import get_correct_switch
 from utils import eprint, create_r2_byte_pattern, sync_r2_output
 import r2pipe
 
@@ -35,6 +35,7 @@ def extract_opcode_data(exe_file):
 
     p = create_r2_byte_pattern(ON_RECEIVE_PACKET_SIG)
     target = r2.cmd(f"/x {p}").split()[0]  # Find byte pattern
+    packet_handler_ea = int(target, 16)
 
     r2.cmd(f"s {target}")  # Seek to target
 
@@ -56,7 +57,9 @@ def extract_opcode_data(exe_file):
 
     ## STEP 4: Process data
     opcodes_db = dict()
-    packet_handler_switch = get_longest_switch(switch_cases)
+    switch_ea, packet_handler_switch = get_correct_switch(packet_handler_ea, switch_cases)
+    eprint(f"  Found switch at {switch_ea}")
+
     vtable_offset = 0x10
     for data in packet_handler_switch.values():
         if len(data["opcodes"]) > 10: