import struct, sys

MODULE = '/tmp/amdgpu_installed.ko'
OUTPUT = '/tmp/amdgpu_repatched.ko'

# Current validation values (from patch_validation.py): MIN=350, MAX=2230
# lea -(2231) = 0xFFFFF749  -> 49 F7 FF FF
# cmp -(2230-350+2) = -1882 = 0xFFFFF8A6 -> A6 F8 FF FF

# New validation values: MIN=350, MAX=2500
# lea -(2501) = 0xFFFFF63B -> 3B F6 FF FF
# cmp -(2500-350+2) = -2152 = 0xFFFFF798 -> 98 F7 FF FF

old_lea = bytes.fromhex('49f7ffff')
new_lea = bytes.fromhex('3bf6ffff')
old_cmp = bytes.fromhex('a6f8ffff')
new_cmp = bytes.fromhex('98f7ffff')

import shutil
shutil.copy2(MODULE, OUTPUT)

with open(OUTPUT, 'rb') as f:
    data = f.read()

print(f'Module size: {len(data)} bytes')
print()

# Find and patch all validation lea/cmp pairs
# Expected offsets: around 0x330186, 0x33018D, 0x3301F9, 0x3301FE
# But these may have shifted if the module binary changed

patches_applied = 0

# Search for old lea displacement in the cyan_skillfish function region (0x330000-0x330400)
for search_start in range(0x330000, 0x330400):
    if data[search_start:search_start+4] == old_lea:
        # Check if there's the cmp value nearby (within 20 bytes)
        for j in range(search_start+4, search_start+24):
            if data[j:j+4] == old_cmp:
                print(f'  Found validation pair: lea@0x{search_start:x}, cmp@0x{j:x}')
                # Patch both
                data = bytearray(data)
                data[search_start:search_start+4] = new_lea
                data[j:j+4] = new_cmp
                patches_applied += 2
                print(f'    Patched lea: {old_lea.hex()} -> {new_lea.hex()}')
                print(f'    Patched cmp: {old_cmp.hex()} -> {new_cmp.hex()}')
                break

if patches_applied == 0:
    print('WARNING: No validation patches found at expected offsets!')
    print('Searching wider range...')
    # Wider search
    pos = 0
    while True:
        idx = bytes(data).find(old_lea, pos)
        if idx == -1: break
        print(f'  Found old_lea at 0x{idx:x}')
        for j in range(idx+4, min(idx+24, len(data)-4)):
            if data[j:j+4] == old_cmp:
                print(f'  Found matching cmp at 0x{j:x}')
                break
        pos = idx + 1

with open(OUTPUT, 'wb') as f:
    f.write(data)

print(f'\nDone. Applied {patches_applied} patches to {OUTPUT}')