Finds location of number in Diggles/Wiggles.exe that sets the number of pixels for mixing ground materials

find_patch_place.py

# For commandline argv
import sys
# For PE parsing
import pefile

# Convert value to 4 bytes
def val_to_b(value):
    return value.to_bytes(4, byteorder="little")

## Constants to search
STRING_LITERAL_A = b"ground/Msk_%c%c%c%c01.tga" # 1 link
STRING_LITERAL_B = b"ground/Gnd_%c%c%c.tga"     # 2 links
STRING_LITERAL_C = b"ground/Gnd_%c%02d.tga"     # 1 link
VALUE_TO_REPLACE = val_to_b(4096)

## Terms
# RAW = raw file address, where data stored in .exe file
# RVA = relative virtual address, where data stored when loads in RAM
# VA  = virtual address, where app think data stored when running
##

## How to convert addresses
# VA = RVA + ImageBase
# RVA = RAW - SectionRAW + SectionRVA
##

# Return relative virtual adresses of section start and end
def get_section_RVA(pe, name):
    if IS_OLD and name in [".text", ".rdata"]:
        name = ".code"
    for section in pe.sections:
        if section.Name.decode("UTF-8").rstrip("\0") == name:
            start = section.VirtualAddress
            end = start + section.SizeOfRawData
            return start, end
    raise KeyError("No section with name %s" % name)

# Get all virtual addresses of value in section
def find_all_VA_in_section(pe, value, name):
    start, end = get_section_RVA(pe, name)
    src = pe.get_memory_mapped_image()
    i = src.find(value, start, end - len(value))
    while i != -1:
        yield i + pe.OPTIONAL_HEADER.ImageBase
        i = src.find(value, i + 1, end - len(value))

# Return address from b that closest to every in a
def get_nearest_address(addr_a, addr_b):
    tmp = [(- sum([(a - b) ** 2 for a in addr_a]), b) for b in addr_b]
    return sorted(tmp, reverse=True)[0][1]

# Convert virtual adress to raw file offset
def VA_to_RAW(pe, addr):
    for section in pe.sections:
        start = section.VirtualAddress + pe.OPTIONAL_HEADER.ImageBase
        end = start + section.SizeOfRawData

        if start <= addr <= end:
            print(section.Name.decode("UTF-8").rstrip("\0"))
            return addr - start + section.PointerToRawData
    raise IndexError("Can't find %s in any sections" % hex(addr))

## What we need to do
# 0. Load PE (.exe) file and prepare section info
# 1. Find all three string literals (in section .rdata)
# 2. Find XREFs to this literals from code (in section .text)
# 3. Drop excess XREF to literal B, so there three XREFs at all
# 4. Find value 4096 in code (in section .text)
# 5. Find distance from every place of 4096 to every XREF to literals
# 6. Choose place with minimal distance as result
##

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: %s Diggles.exe" % sys.argv[0])
        sys.exit(-1)

    # 0. Load PE (.exe) file and prepare section info
    pe = pefile.PE(name=sys.argv[1], fast_load=True)

    # Check for UPX packing and old versions (only one section)
    IS_OLD = False
    for section in pe.sections:
        if section.Name.decode("UTF-8")[:3] == "UPX":
            print("\"%s\" is packed using UPX; \
unpack it with https://upx.github.io" % sys.argv[1])
            sys.exit(-1)

        if section.Name.decode("UTF-8")[:5] == ".code":
            print("\"%s\" is old version of game (UPX unpacked)\n" %
                  sys.argv[1])
            IS_OLD = True

    # 1. Find all three string literals (in section .rdata)
    literal_a = list(find_all_VA_in_section(pe, STRING_LITERAL_A, ".rdata"))
    assert len(literal_a) == 1, "Found %d instead of 1" % len(literal_a)
    literal_a = literal_a[0]
    print("Found string A \"%s\" at VA %s" %
          (STRING_LITERAL_A.decode("UTF-8"), hex(literal_a)))

    literal_b = list(find_all_VA_in_section(pe, STRING_LITERAL_B, ".rdata"))
    assert len(literal_b) == 1, "Found %d instead of 1" % len(literal_b)
    literal_b = literal_b[0]
    print("Found string B \"%s\" at VA %s" %
          (STRING_LITERAL_B.decode("UTF-8"), hex(literal_b)))

    literal_c = list(find_all_VA_in_section(pe, STRING_LITERAL_C, ".rdata"))
    assert len(literal_c) == 1, "Found %d instead of 1" % len(literal_c)
    literal_c = literal_c[0]
    print("Found string C \"%s\" at VA %s" %
          (STRING_LITERAL_C.decode("UTF-8"), hex(literal_c)))

    # 2. Find XREFs to this literals from code (in section .text)
    xref_a = list(find_all_VA_in_section(pe, val_to_b(literal_a), ".text"))
    assert len(xref_a) == 1, "Found %d instead of 1" % len(xref_a)
    xref_a = xref_a[0]
    print("Found XREF to string A at VA %s" % hex(xref_a))

    xref_b = list(find_all_VA_in_section(pe, val_to_b(literal_b), ".text"))
    assert len(xref_b) == 2, "Found %d instead of 2" % len(xref_b)
    print("Found XREF to string B at VA %s and %s" %
          (hex(xref_b[0]), hex(xref_b[1])))

    xref_c = list(find_all_VA_in_section(pe, val_to_b(literal_c), ".text"))
    assert len(xref_c) == 1, "Found %d instead of 1" % len(xref_c)
    xref_c = xref_c[0]
    print("Found XREF to string C at VA %s" % hex(xref_c))

    # 3. Drop excess XREF to literal B, so there three XREFs at all
    xref_b = get_nearest_address([xref_a, xref_c], xref_b)
    print("Keep XREF to string B from VA %s" % hex(xref_b))

    # 4. Find value 4096 in code (in section .text)
    pos_value = list(find_all_VA_in_section(pe, VALUE_TO_REPLACE, ".text"))
    assert len(pos_value) > 0, "Can't find %s in .text" % VALUE_TO_REPLACE
    print("Found %d places with %s" % (len(pos_value), VALUE_TO_REPLACE))

    # 5. Find distance from every place of 4096 to every XREF to literals
    # 6. Choose place with minimal distance as result
    pos_value = get_nearest_address([xref_a, xref_b, xref_c], pos_value)
    print("Finally chose value place as VA %s" % hex(pos_value))

    print("\nReplace 4 bytes at position %s (%d) in file \"%s\"" %
          (hex(pos_value), pos_value, sys.argv[1]))