#jmp_table_addr = 0x6ea1a0
#n = 0x6a
#index_offset = 0x3c8
#addr_array_stack_addr = 0x360

#jmp_table_addr = 0x6ea500
#n = 0x61
#index_offset = 0x378
#addr_array_stack_addr = 0x310

#jmp_table_addr = 0x6ea820
#n = 0x56
#index_offset = 0x338
#addr_array_stack_addr = 0x2c0

#jmp_table_addr = 0x6eaae0
#n = 0x62
#index_offset = 0x3e8
#addr_array_stack_addr = 0x350

#jmp_table_addr = 0x6eae00
#n = 0x68
#index_offset = 0x3e8
#addr_array_stack_addr = 0x350

#jmp_table_addr = 0x6eb140
#n = 0x55
#index_offset = 0x348
#addr_array_stack_addr = 0x2b0

#jmp_table_addr = 0x6eb400
#n = 0x65
#index_offset = 0x3d8
#addr_array_stack_addr = 0x330

#jmp_table_addr = 0x6eb740
#n = 0x63
#index_offset = 0x3f8
#addr_array_stack_addr = 0x320

jmp_table_addr = 0x6eba60
n = 0x5a
index_offset = 0x388
addr_array_stack_addr = 0x2e0

#jmp_table_addr = 0x6ebd40
#n = 0x6b
#index_offset = 0x3c8
#addr_array_stack_addr = 0x360

def is_in_ranges(addr, addr_ranges):
    for addr_range in addr_ranges:
        if addr_range[0] <= addr <= addr_range[1]:
            return True

    return False

valid_addr_range = []
refs = bv.get_code_refs(jmp_table_addr)
base_func = bv.get_functions_containing(refs[0].address)[0]
valid_addr_range.append((base_func.lowest_address, base_func.highest_address))

arch = bv.arch
br = BinaryReader(bv)
br.seek(jmp_table_addr)
addrs = []

bv.begin_undo_actions()
for i in range(n):
    addr = br.read64le()
    addrs.append(addr)
    if not bv.get_function_at(addr):
        bv.add_function(addr)
        log_info(f'new function added at 0x{addr:x}')

bv.update_analysis_and_wait()

for addr in addrs:
    func = bv.get_function_at(addr)
    if not func:
        log_warn(f'no function at {addr:#x}')
        raise

    valid_addr_range.append((func.lowest_address, func.highest_address))

search_text = f'mov     qword [rbp-0x{index_offset:x}'
log_info(f'searching for: {search_text}')
count = 0

for result in bv.find_all_text(bv.start, bv.end, search_text):
    addr, s, line = result
    if not is_in_ranges(addr, valid_addr_range):
        continue

    log_info(f'{addr: #x}')
    func = bv.get_functions_containing(addr)[0]
    llil_func = func.llil
    instr = func.get_llil_at(addr)
    idx = instr.instr_index
    next_llil_instr = llil_func[idx + 1]
    if next_llil_instr.operation == LowLevelILOperation.LLIL_GOTO:
        check_il = llil_func[next_llil_instr.dest].instr_index + 1
    else:
        check_il = idx + 2

    check_il_instr = llil_func[check_il]
    try:
        offset = check_il_instr.src.src.right.constant
    except:
        log_warn(f'fail to process 0x{addr:x}, {s}')
        continue

    if offset != -addr_array_stack_addr:
        # This is NOT what we are interested!
        continue

    if instr.operation == LowLevelILOperation.LLIL_STORE:
        if instr.src.operation == LowLevelILOperation.LLIL_CONST:
            const = instr.src.constant

    # log_info(f'const: {const}')
    try:
        new_addr = addrs[const]
    except:
        new_addr = 0x0
        log_warn(f'jumping to 0x0 at {addr: #x}')

    code = arch.assemble(f'jmp {new_addr: #x}', addr)
    log_info(f'{addr: #x} ==> {new_addr: #x}')
    bv.write(addr, code)
    count += 1

log_info(f'{count} jumps processed')
for addr in addrs:
    # undefine all of these short functions, so that it becomes one single function
    func = bv.get_function_at(addr)
    if func:
        bv.remove_user_function(func)

bv.commit_undo_actions()
bv.update_analysis_and_wait()

log_info(f'done')

