blockland-DLL-tools/BL-patcher.py

import pefile
import copy
from hashlib import sha256

'''
Rip out the DOS stub:

dd if=Blockland.exe bs=1 count=296 skip=64 of=dos_stub.txt

'''
dll_target_name = "RedBlocklandLoader.dll"

print(" **************************************************************************")
print(" *                                                                        *")
print(" *                           Blockland-patcher                            *")
print(" *                                                                        *")
print(" **************************************************************************")
print()

#one final check: ignore a file that already has newimp

exe_path = "Blockland.r2001.std.exe"
exe_path = "Blockland.r1986.std.exe"
exe_path = "Blockland.r2012.std.exe"
exe_path = "Blockland.exe"

print(' - Input file: '+ exe_path)
print(' - Interpreting EXE file...')
pe = pefile.PE(exe_path)
added_size_counter = 0


print(' - Re-aligning all sections...')

for section in pe.sections:
	virt_size = section.Misc_VirtualSize
	raw_size = section.SizeOfRawData
	phys_addr = section.Misc_PhysicalAddress

	if virt_size % 0x200 != 0 and virt_size < raw_size:
		old_size = phys_addr
		new_size = (int(old_size / 0x200) + 1) * 0x200 #round up to nearest 0x200
		section.Misc_PhysicalAddress = new_size
		#section.SizeOfRawData = new_size

		added_size_counter += new_size-old_size
		name_strip = section.Name.decode().strip()
		print("    Section boundary does not align to page size. Adjusting section \""+ name_strip +"\":")
		print("        From: "+ str(old_size) +" ("+ hex(old_size) +")")
		print("          To: "+ str(new_size) +" ("+ hex(new_size) +")")
print()


newSection = copy.copy(pe.sections[0])

newSection.Name=b".newimp"
newSection.Misc = newSection.Misc_PhysicalAddress = newSection.Misc_VirtualSize = 0x1000
newSection.VirtualAddress = pe.sections[-1].SizeOfRawData+pe.sections[-1].VirtualAddress
newSection.SizeOfRawData = 0x1000
newSection.PointerToRawData = pe.sections[-1].SizeOfRawData+pe.sections[-1].PointerToRawData
newSection.Characteristics = 0xC0000040

newSection.set_file_offset(pe.sections[-1].__file_offset__ + newSection.sizeof())

pe.OPTIONAL_HEADER.SizeOfImage += newSection.SizeOfRawData
pe.FILE_HEADER.NumberOfSections += 1

pe.sections.append(newSection)
pe.__structures__.append(newSection)

print(" - Finding the import address table...")
found_import_addr = False
for entry in pe.OPTIONAL_HEADER.DATA_DIRECTORY:
	if entry.name == "IMAGE_DIRECTORY_ENTRY_IMPORT":
		found_import_addr = True

		print(" - Modifying ImportAddressTable and SizeOfImportAddressTable to point to .newimp section:")
		print("    Old size: "+  hex(entry.Size))
		print("    Old virtual address: "+ hex(entry.VirtualAddress))

		entry.Size += 20
		entry.VirtualAddress = newSection.VirtualAddress;

		print("    New size: "+  hex(entry.Size))
		print("    New virtual address: "+ hex(entry.VirtualAddress))

		break

if not found_import_addr:
	print(" - Unable to find ImportAddressTable.")
	quit()

print()

new_import_table=bytearray(0)

print(' - Reading original import table...')

newSection_current_offset = newSection.VirtualAddress

for entry in pe.DIRECTORY_ENTRY_IMPORT:
	dll_name=entry.dll.decode('utf-8')

	print('    Found '+ dll_name)
	
	struct=entry.struct

	data = [
		struct.Characteristics, #AKA pointer
		struct.TimeDateStamp,
		struct.ForwarderChain,
		struct.Name,
		struct.FirstThunk
	]

	for var in data:
		new_import_table.extend(var.to_bytes(4, byteorder='little'))
		newSection_current_offset += 4

total_imports = len(pe.DIRECTORY_ENTRY_IMPORT)
print('    ('+ str(total_imports) +' imports)')
print()


print(' - Constructing new import table entry...')

dll_target_name_term=dll_target_name +"\0"
#if len(dll_target_name_term) % 2 == 1:
#	dll_target_name_term += '\0'

#Each address skips 20 bytes over the current entry and 20 over the terminal empty entry

#import table address
addr=newSection_current_offset+40+len(dll_target_name_term)+8
new_import_table.extend(addr.to_bytes(4, byteorder='little'))

#timestamp
new_import_table.extend(bytearray.fromhex("00 00 00 00"))

#forwarder chain
new_import_table.extend(bytearray.fromhex("00 00 00 00"))

#address of the DLL's name string
addr=newSection_current_offset+40+0
new_import_table.extend(addr.to_bytes(4, byteorder='little'))

#first thunk
addr=newSection_current_offset+40+len(dll_target_name_term)
new_import_table.extend(addr.to_bytes(4, byteorder='little'))

newSection_current_offset += 20


#empty import table entry marks end of import table
new_import_table.extend(bytearray(20))
newSection_current_offset += 20



print(' - Constructing pointers for '+ dll_target_name +'...')

new_loader_lookup_table=bytearray(0)

#byte 0: DLL name
new_loader_lookup_table.extend(str.encode(dll_target_name_term))
#By spec, aligned to even number of bytes, but doesn't seem to matter


newSection_current_offset += len(dll_target_name_term)

#byte 20: first entry of import address table
addr=newSection_current_offset+16
new_loader_lookup_table.extend(addr.to_bytes(4, byteorder='little'))

#byte 24: padding
new_loader_lookup_table.extend(bytearray.fromhex("00 00 00 00"))

#byte 28: first entry of import lookup table
addr=newSection_current_offset+16
new_loader_lookup_table.extend(addr.to_bytes(4, byteorder='little'))

#byte 32: padding
new_loader_lookup_table.extend(bytearray.fromhex("00 00 00 00"))

#byte 36: export name pointer table is size 0
new_loader_lookup_table.extend(bytearray.fromhex("00 00"))

#byte 38: the function referenced in the DLL
new_loader_lookup_table.extend(b"RedBlocklandLoader\0")


section_data = new_import_table + new_loader_lookup_table;
added_size_counter += len(section_data)


if added_size_counter % 4096 == 0:
	pe.OPTIONAL_HEADER.SizeOfInitializedData += added_size_counter
else:
	pe.OPTIONAL_HEADER.SizeOfInitializedData += (added_size_counter // 4096 + 1) * 4096


print(' - Writing new EXE with modified headers...')
new_exe_path = "Blockland.patched.exe"
pe.write(new_exe_path)



print(' - Re-writing EXE to append content of .newimp section...')




null_size = 4096 - len(section_data)
section_data += bytearray(null_size)

section_data_hash = sha256(section_data).hexdigest()


#DOS stub doesn't actually change, but it is not preserved when patching BL with another program (perhaps StudPE)
#The code to change the stub is kept here in case I want a 100%-matching binary
#with open('dos_stub.txt', 'r') as file:
#	dos_stub = file.read()

#with open('Blockland.exe', 'r') as file:
#	exe_data = file.read()

#with open("Blockland.exe", "w") as file:
#	file.write(exe_data[:64] + dos_stub + exe_data[360:] + section_data)


with open(new_exe_path, 'rb') as file:
	exe_data = file.read()

with open(new_exe_path, "wb") as file:
	file.write(exe_data + section_data)



pe = pefile.PE(new_exe_path)



print(" - Checking .newimp section checksum")

refound_section = False
for section in pe.sections:
	if section.Name.decode() != '.newimp\0':
		continue

	check_section_data_hash = sha256(section.get_data()).hexdigest()
	if check_section_data_hash == section_data_hash:
		print("    Hash match.")
	else:
		print("    Hash mismatch.")
		quit()

	refound_section = True
	break

if not refound_section:
	print("    Could not find section.")
	quit()


print()
print(" - Checking presence of new "+ dll_target_name +":")

dll_found = False
dll_count = 0

for entry in pe.DIRECTORY_ENTRY_IMPORT:
	dll_name=entry.dll.decode('utf-8')

	dll_count += 1

	if not dll_name == dll_target_name:
		continue

	print("    Found! "+ dll_name +" functions:")

	for func in entry.imports:
		print("        %s at 0x%08x" % (func.name.decode('utf-8'), func.address))

	dll_found = True
	break

if not dll_found:
	print("    Could not find reference to "+ dll_target_name +" and its imports.")
	quit()

print("    The EXE contains "+ str(dll_count) +" imports.")

print()
print(" **************************************************************************")
print(" *                                                                        *")
print(" *                 The patched EXE has passed all checks.                 *")
print(" *                                                                        *")
print(" **************************************************************************")
print()