"""
:: File Name:
wow_script_functions.py

:: Author:
Jadd - https://ntoskr.nl/

:: Target:
World of Warcraft x64 - Retail Client 8.0+
World of Warcraft x64 - Classic Client 1.13+

:: Description:
Renames Lua scripting functions within the game client, more specifically
  those referencing one of four common "FrameScript" functions. Generates a
  function renaming script by outputting the results to lua_functions.py in
  the IDB's directory.

:: License:
Free to use under MIT License.
"""


# o_reg values.
REG_AX  = 0;  REG_CX  = 1;  REG_DX  = 2;  REG_BX  = 3;
REG_SP  = 4;  REG_BP  = 5;  REG_SI  = 6;  REG_DI  = 7;
REG_R8  = 8;  REG_R9  = 9;  REG_R10 = 10; REG_R11 = 11;
REG_R12 = 12; REG_R13 = 13; REG_R14 = 14; REG_R15 = 15;


def find_prev_op(ea, operand):
    if ea == BADADDR:
        return BADADDR
    
    search_end = first_func_chunk(ea)

    while ea != BADADDR:
        ea = prev_head(ea)
        if ea < search_end:
            break
        if print_insn_mnem(ea) == operand:
            return ea

    return BADADDR


def find_next_op(ea, operand):
    if ea == BADADDR:
        return BADADDR
    
    search_end = find_func_end(ea)

    while ea != BADADDR:
        ea = next_head(ea)
        if ea >= search_end:
            break
        if print_insn_mnem(ea) == operand:
            return ea

    return BADADDR


def search_framescript_function(name):
    if name == 'FrameScript_RegisterFunction':
        ea = FindBinary(0, SEARCH_DOWN, '48 8B D9 48 8B CF 45 33 C0 E8 ?? ?? ?? ?? 48 8B D3 48 8B CF E8 ?? ?? ?? ?? BA FE FF FF FF')
        assert(ea != BADADDR)
        ea = first_func_chunk(ea)
        return ea

    if name == 'FrameScript_RegisterGlobalFunction':
        ea = FindBinary(0, SEARCH_DOWN, '49 8B C0 48 8B DA 48 8B F9 48 8B D0 48 8B 09 45 33 C0 E8 ?? ?? ?? ??')
        assert(ea != BADADDR)
        ea = first_func_chunk(ea)
        return ea

    if name == 'FrameScript_RegisterTableFunction':
        ea = FindBinary(0, SEARCH_DOWN, '41 B8 EE D8 FF FF 4C 8B F1 E8 ?? ?? ?? ?? 49 8B 0E E8 ?? ?? ?? ?? 49 8B 0E')
        assert(ea != BADADDR)
        ea = first_func_chunk(ea)
        return ea

    if name == 'FrameScript_RegisterFunctionNamespaceWithCount':
        ea = FindBinary(0, SEARCH_DOWN, 'BA EE D8 FF FF 48 8B CB 48 8B 5C 24 ?? 48 8B 6C 24 ?? 48 8B 74 24 ??')
        assert(ea != BADADDR)
        ea = first_func_chunk(ea)
        return ea


def read_lua_functions(framescript_function):
    results = []
    reference = get_next_fcref_to(framescript_function, 0)
    
    while reference != BADADDR:
        if print_insn_mnem(reference) == "jmp":
            ea = prev_head(reference)
            rcx = get_operand_value(ea, 1)
            ea = prev_head(ea)
            rdx = get_operand_value(ea, 1)
            function_name = get_strlit_contents(get_qword(rcx), -1, ASCSTR_C)
            function_addr = get_qword(rdx) - get_first_seg()
            results.append({'name': function_name, 'address': function_addr})
            reference = get_next_fcref_to(framescript_function, reference)
            continue

        compare = reference
        ea = reference

        registers = {}
        function_count = 1
        mnemonic = ''
        
        # Check how it determines the function table length.
        while mnemonic != 'cmp' and mnemonic != 'sub':
            compare = next_head(compare)
            mnemonic = print_insn_mnem(compare)
        
        # Possible operations:
        # cmp rbx, ... -> pointer comparison
        # cmp rdi, ... -> counter comparison
        # sub rdi, 1 -> countdown
        comparison = get_operand_value(compare, 0)
        countdown = mnemonic == 'sub'
        
        if comparison != REG_BX and comparison != REG_DI:
            raise Exception('Invalid comparison at 0x{:016X}'.format(reference))
        
        if comparison == REG_DI and not countdown:
            registers[REG_DI] = get_operand_value(compare, 1)

        while REG_BX not in registers or REG_DI not in registers:
            mnemonic = ''
            while mnemonic != 'lea' and mnemonic != 'mov':
                ea = prev_head(ea)
                mnemonic = print_insn_mnem(ea)

            register = get_operand_value(ea, 0)
            if register not in registers:
                registers[register] = get_operand_value(ea, 1)
        
        if comparison == REG_BX:
            function_count = (registers[REG_DI] - registers[REG_BX]) / 0x10
        else:
            function_count = registers[REG_DI]

        for i in xrange(0, function_count):
            entry = registers[REG_BX] + (i * 0x10)
            function_name = get_strlit_contents(get_qword(entry), -1, ASCSTR_C)
            function_addr = get_qword(entry + 8) - get_first_seg()
            results.append({'name': function_name, 'address': function_addr})
        
        reference = get_next_fcref_to(framescript_function, reference)
        
    return results


def read_lua_global_functions(framescript_function):
    results = []
    reference = get_next_fcref_to(framescript_function, 0)
    
    while reference != BADADDR:
        registers = {}
        ea = reference
        
        while REG_DX not in registers or REG_R8 not in registers:
            ea = find_prev_op(ea, 'lea')
            registers[get_operand_value(ea, 0)] = get_operand_value(ea, 1)
        
        function_name = get_strlit_contents(registers[REG_DX], -1, ASCSTR_C)
        function_addr = registers[REG_R8] - get_first_seg()
        results.append({'name': function_name, 'address': function_addr})
        
        reference = get_next_fcref_to(framescript_function, reference)
        
    return results


def read_lua_namespace_functions(framescript_function):
    results = []
    reference = get_next_fcref_to(framescript_function, 0)
    
    while reference != BADADDR:
        registers = {}
        ea = reference
        
        while REG_DX not in registers or REG_R8 not in registers or REG_R9 not in registers:
            ea = find_prev_op(ea, 'lea')
            registers[get_operand_value(ea, 0)] = get_operand_value(ea, 1)
        
        namespace = get_strlit_contents(registers[REG_DX], -1, ASCSTR_C)
        function  = get_strlit_contents(registers[REG_R8], -1, ASCSTR_C)

        function_name = '{}.{}'.format(namespace, function)
        function_addr = registers[REG_R9] - get_first_seg()
        results.append({'name': function_name, 'address': function_addr})
        
        reference = get_next_fcref_to(framescript_function, reference)
        
    return results


def read_lua_namespace_with_count_functions(framescript_function):
    results = []
    reference = get_next_fcref_to(framescript_function, 0)
    
    while reference != BADADDR:
        registers = {}
        ea = reference

        while REG_CX not in registers or REG_DX not in registers or REG_R8 not in registers:
            mnemonic = ''
            while mnemonic != 'lea' and mnemonic != 'mov':
                ea = prev_head(ea)
                mnemonic = print_insn_mnem(ea)
            registers[get_operand_value(ea, 0)] = get_operand_value(ea, 1)
        
        namespace = get_strlit_contents(registers[REG_R8], -1, ASCSTR_C)

        for i in xrange(0, registers[REG_DX]):
            function = get_strlit_contents(get_qword(registers[REG_CX] + i * 16), -1, ASCSTR_C)
            function_name = '{}.{}'.format(namespace, function)
            function_addr = get_qword(registers[REG_CX] + 8 + i * 16) - get_first_seg()
            results.append({'name': function_name, 'address': function_addr})
        
        reference = get_next_fcref_to(framescript_function, reference)
        
    return results


def log_output(file, definitions):
    for function in definitions:
        file.write('MakeNameEx(baseAddr + 0x{address:016X}, "Script_{name}", SN_NOWARN)\n'.format(**function))
        file.write('SetType(baseAddr + 0x{address:016X}, "signed int __fastcall Script_{name}(void *L);")\n'.format(**function))
        print('0x{address:016X} Script_{name}'.format(**function))


def define_lua_functions(definitions):
    for function in definitions:
        MakeNameEx(get_first_seg() + function['address'], 'Script_{name}'.format(**function), SN_NOWARN)
        SetType(get_first_seg() + function['address'], 'signed int __fastcall Script_{name}(void *L);'.format(**function))


def main():
    FrameScript_RegisterFunction = search_framescript_function('FrameScript_RegisterFunction')
    FrameScript_RegisterGlobalFunction = search_framescript_function('FrameScript_RegisterGlobalFunction')
    FrameScript_RegisterTableFunction = search_framescript_function('FrameScript_RegisterTableFunction')
    FrameScript_RegisterFunctionNamespaceWithCount = search_framescript_function('FrameScript_RegisterFunctionNamespaceWithCount')

    MakeNameEx(FrameScript_RegisterFunction, 'FrameScript_RegisterFunction', SN_NOWARN)
    SetType(FrameScript_RegisterFunction, 'void __fastcall FrameScript_RegisterFunction(const char *name, signed int (__fastcall *function)(void *L))')

    MakeNameEx(FrameScript_RegisterGlobalFunction, 'FrameScript_RegisterGlobalFunction', SN_NOWARN)
    SetType(FrameScript_RegisterGlobalFunction, 'void __fastcall FrameScript_RegisterGlobalFunction(void* L, const char *name, signed int (__fastcall *function)(void *L))')

    MakeNameEx(FrameScript_RegisterTableFunction, 'FrameScript_RegisterTableFunction', SN_NOWARN)
    SetType(FrameScript_RegisterTableFunction, 'void __fastcall FrameScript_RegisterTableFunction(void* L, const char *table, const char *name, signed int (__fastcall *function)(void *L))')

    MakeNameEx(FrameScript_RegisterFunctionNamespaceWithCount, 'FrameScript_RegisterFunctionNamespaceWithCount', SN_NOWARN)
    SetType(FrameScript_RegisterFunctionNamespaceWithCount, 'void __fastcall FrameScript_RegisterFunctionNamespaceWithCount(__int64* functionTable, int functionCount, const char *table')

    lua_functions = read_lua_functions(FrameScript_RegisterFunction)
    define_lua_functions(lua_functions)

    lua_global_functions = read_lua_global_functions(FrameScript_RegisterGlobalFunction)
    define_lua_functions(lua_global_functions)

    lua_namespace_functions = read_lua_namespace_functions(FrameScript_RegisterTableFunction)
    define_lua_functions(lua_namespace_functions)

    lua_namespace_with_count_functions = read_lua_namespace_with_count_functions(FrameScript_RegisterFunctionNamespaceWithCount)
    define_lua_functions(lua_namespace_with_count_functions)

    with open('lua_functions.py', 'w') as log:
        log.write('baseAddr = idc.get_first_seg()\n\n')
        
        log.write('# Global Lua Functions #\n')
        log_output(log, lua_functions)
        log_output(log, lua_global_functions)

        log.write('\n')
        log.write('# Namespace Lua Functions #\n')
        log_output(log, lua_namespace_functions)
        log_output(log, lua_namespace_with_count_functions)


if __name__ == '__main__':
    main()