#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-only
# Copyright David Heidelberg.
"""
Patch SHIFT 6MQ wlanmdsp.mbn to fix NULL function pointer crash.

Bug: WLAN.HL.2.0-03435-QCAHLSWMTPLZ-1 (SHIFT 6MQ stock wlanmdsp.mbn) crashes
at PC=0xb001860c with a NULL function pointer dereference (callr r3, where r3=0).

Root cause: Internal command handler 0x0a enters a per-radio iteration loop
that calls through a function pointer at BSS address 0xb05fddf0. This pointer
is never initialized (BSS is zero-filled), so the firmware crashes with:

    err_qdi.c:456:EX:wlan_process:1:WLAN RT:2080:PC=b001860c

The crash path is taken when *(context+4) == 0, which indicates the vdev
context is in an uninitialized state. The normal path (*(context+4) != 0)
jumps to a different handler at 0xb0018648.

Fix: Replace the first instruction in the crash path with an unconditional
jump to the function's clean exit epilogue. This makes the handler safely
return when the context field is zero, instead of entering the buggy loop.

    Before: { r0 = memw(gp+#0x154) }   @ 0xb00185e0
    After:  { jump 0xb0018ae4 }         @ 0xb00185e0  (-> epilogue)

The patched firmware functions identically to the OP6 wlanmdsp.mbn
(WLAN.HL.2.0-01387) which handles this case correctly by having different
code at this address entirely.
"""

import hashlib
import shutil
import struct
import sys
import os

PATCH_OFFSET = 0x285e0      # File offset (PC 0xb00185e0)
ORIGINAL_INSTR = 0x4980caa0  # r0 = memw(gp+#0x154)
PATCHED_INSTR  = 0x5800c282  # jump 0xb0018ae4

HASH_TABLE_SEG2_OFFSET = 0x1068  # SHA-256 hash of text segment in hash table
TEXT_SEGMENT_OFFSET = 0x10000
TEXT_SEGMENT_SIZE = 0x28c398


def patch_firmware(input_path, output_path=None):
    if output_path is None:
        base, ext = os.path.splitext(input_path)
        output_path = f"{base}_patched{ext}"

    shutil.copy2(input_path, output_path)

    with open(output_path, 'r+b') as f:
        # Verify original instruction
        f.seek(PATCH_OFFSET)
        original = struct.unpack('<I', f.read(4))[0]
        if original != ORIGINAL_INSTR:
            if original == PATCHED_INSTR:
                print(f"Already patched: {input_path}")
                return output_path
            print(f"ERROR: Unexpected instruction at offset {PATCH_OFFSET:#x}: "
                  f"{original:#010x} (expected {ORIGINAL_INSTR:#010x})")
            sys.exit(1)

        # Apply code patch
        f.seek(PATCH_OFFSET)
        f.write(struct.pack('<I', PATCHED_INSTR))

        # Update hash table for text segment
        f.seek(TEXT_SEGMENT_OFFSET)
        text_data = f.read(TEXT_SEGMENT_SIZE)
        new_hash = hashlib.sha256(text_data).digest()

        f.seek(HASH_TABLE_SEG2_OFFSET)
        f.write(new_hash)

    print(f"Patched: {output_path}")
    print(f"  Code:  {ORIGINAL_INSTR:#010x} -> {PATCHED_INSTR:#010x} @ offset {PATCH_OFFSET:#x}")
    print( "  Hash:  Updated text segment SHA-256 in hash table")
    return output_path


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <wlanmdsp.mbn> [output.mbn]")
        sys.exit(1)

    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else None
    patch_firmware(input_file, output_file)