script.py

from operator import le
import re

file = open('PLAY.mal_', 'rb')
buffer = file.read()
file.close()

# main case
regex = b'\x06\x81\xC4....\x83\xC4.\xE8....'
count = 0


def patchedBytes(buffer, changes):

    for change in changes:
        offset = change[0]
        length = change[1]
        value = change[2]

        if len(value) != length:
            return
        buffer = [each for each in buffer]
        value = [each for each in value]

        for i in range(length):
            buffer[i + offset] = value[i]

    result = b''

    for each in buffer:
        result += int.to_bytes(each, 1, 'little')
    return result


def toIDAAddress(address):
    return address - 0x400 + 0x401000


def toBufferAddress(address):
    return address + 0x400 - 0x401000


fix_dict = {}
changes = []
for match in re.finditer(regex, buffer):
    count += 1
    new_esp = match.start() + 15
    match_bytes = match.group()
    call_address = int.from_bytes(
        match_bytes[match_bytes.find(b'\xE8') + 1:], "little") + new_esp

    addition_val = int(buffer[call_address + 3])
    jump_address = new_esp + addition_val

    patch_address = new_esp - 5
    changes.append([patch_address, 5, b'\xE9' +
                   int.to_bytes(addition_val, 1, 'little') + b'\x00\x00\x00'])

buffer = patchedBytes(buffer, changes)

# edge_case 1

regex = b'\x83\xC4.\x83\xC4.\xE8....'
count = 0

match = re.findall(regex, buffer)

for match in re.finditer(regex, buffer):
    count += 1
    new_esp = match.start() + 11
    match_bytes = match.group()
    call_address = int.from_bytes(
        match_bytes[match_bytes.find(b'\xE8') + 1:], "little") + new_esp
    print(hex(new_esp), hex(call_address), hex(toIDAAddress(match.start())))
    addition_val = int(buffer[call_address + 3])
    jump_address = new_esp + addition_val

    patch_address = new_esp - 5
    changes.append([patch_address, 5, b'\xE9' +
                   int.to_bytes(addition_val, 1, 'little') + b'\x00\x00\x00'])

buffer = patchedBytes(buffer, changes)

# edge_case 2

regex = b'\x83\xC4.\xE8'
count = 0

match = re.findall(regex, buffer)

for match in re.finditer(regex, buffer):
    count += 1
    if toIDAAddress(match.start()) <= 0x410000:
        continue
    new_esp = match.start() + 8
    match_bytes = match.group()
    call_address = int.from_bytes(
        buffer[match.start() + 4:match.start() + 8], "little") + new_esp
    addition_val = int(buffer[call_address + 3])
    jump_address = new_esp + addition_val

    patch_address = new_esp - 5
    changes.append([patch_address, 5, b'\xE9' +
                   int.to_bytes(addition_val, 1, 'little') + b'\x00\x00\x00'])

# write changes
buffer = patchedBytes(buffer, changes)

file = open('patched_PLAY.mal_', 'wb')
file.write(buffer)
file.close()