IDAPython script to re-enable stubbed logs from jubeat (2008)

jubeat1_log_enabler.py

# Jubeat 1 has logging functions implemented using a #define such as:
# #define LOG_MISC(fmt, ...) while(LOGS_ENABLED) { log_body_misc(fmt, __VA_ARGS__); }
# 
# However, because barely any optimisations were applied to the code, these are
# still present in the final binary, even when LOGS_ENABLED was 0. This script
# re-enables these logs.
# 
# An example assembly snippet:
# start:
# xor     eax, eax
# jz      short end
# mov     ecx, [ebp+var_10]
# push    ecx
# push    offset aError ; "error=%d"
# push    offset aModule; "module"
# call    ds:log_body_misc
# add     esp, 0Ch
# jmp     start
# end:
# 
# By simply nopping out the leading xor and jz, and the trailing jmp, we can
# trivially restore the logs.
# This log restoration can be applied to device.dll, jubeat.exe, ledanim.dll,
# libbmsd.dll, netcom2.dll, pkfs.dll, sciunit.dll, and vsnet.dll
#
# For a given file.ext it will create file_logging.ext

import os

funcs = ['log_body_misc','log_body_warning','log_body_info','log_body_fatal']

nops = []

for f in funcs:
    f_addr = get_name_ea(0, f)
    xrefs = XrefsTo(f_addr)
    for ref in xrefs:
        # call    ds:log_body_misc
        call_loc = ref.frm
        if idc.print_insn_mnem(call_loc) != 'call':
            continue
        # add     esp, 10h
        esp_add = next_head(call_loc)
        # jmp     short loc_4027FE
        jump_back = next_head(esp_add)
        if idc.print_insn_mnem(jump_back) != 'jmp':
            continue
        jump_sz = get_item_size(jump_back)

        jump_target = get_operand_value(jump_back, 0)

        # xor     eax, eax
        if idc.print_insn_mnem(jump_target) != 'xor':
            continue
        # jz      short loc_4027ED
        branch_target = next_head(jump_target)
        if idc.print_insn_mnem(branch_target) != 'jz':
            continue
        branch_sz = get_item_size(branch_target)

        nops.extend(range(jump_back, jump_back+jump_sz))
        nops.extend(range(branch_target, branch_target+branch_sz))

src = ida_nalt.get_input_file_path()
base, ext = os.path.splitext(src)
dst = base + '_logging' + ext
with open(src, 'rb') as f:
    binary = bytearray(f.read())

for nop in nops:
    binary[ida_loader.get_fileregion_offset(nop)] = 0x90

with open(dst, 'wb') as f:
    f.write(binary)

print('Done!')