Patch Windows EXE (DLL) with python script

binpatch.py

import os
import shutil
import subprocess
import tempfile

from pwnlib.log import getLogger

log = getLogger(__name__)

def _run(cmd, stdin = None):
	log.debug("%s", subprocess.list2cmdline(cmd))
	try:
		startupinfo = subprocess.STARTUPINFO()
		startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
		proc = subprocess.Popen(
			cmd,
			stdin  = subprocess.PIPE,
			stdout = subprocess.PIPE,
			stderr = subprocess.PIPE,
			universal_newlines = True,
			startupinfo = startupinfo
		)
		stdout, stderr = proc.communicate(stdin)
		exitcode = proc.wait()
	except OSError as e:
		if e.errno == errno.ENOENT:
			log.exception("Could not run %r the program", cmd[0])
		else:
			raise

	if (exitcode, stderr) != (0, ""):
		msg = "There was an error running %r:\n"
		args = cmd,
		if exitcode != 0:
			msg += "It had the exitcode %d.\n"
			args += exitcode,
		if stderr != "":
			msg += "It had this on stdout:\n%s\n"
			args += stderr,
		log.error(msg, *args)
		
	return stdout

def nasm(shellcode):
	code = "BITS 64\n" + shellcode

	tmpdir = tempfile.mkdtemp(prefix = "pwn-nasm-")
	input_file = os.path.join(tmpdir, "shellcode.S")
	cmd  = [
		"nasm",
		input_file
	]
	try:
		open(input_file, "wb").write(code.encode("u8"))
		_run(cmd, code)
		buf = open(os.path.join(tmpdir, "shellcode"), "rb").read()
	except Exception:
		log.exception("An error occurred while compiling shellcode:\n%s" % code)
	shutil.rmtree(tmpdir)
	
	return buf

def pprintBytes(s):
	return " ".join([hex(i)[2:].rjust(2, "0") for i in s])
	
def virtToPhysOffset(virtAddress):
	return virtAddress - 0x78E51000 + 1024
	
def applyPatch(s, patchLocVirt, patch):
	patchLocPhys = virtToPhysOffset(patchLocVirt)
	for i in range(patchLocPhys, patchLocPhys + len(patch)):
		s[i] = patch[i - patchLocPhys]
		
	return s
	
def main():
	firstPatchLoc = 0x78E9EDD0
	secondPatchLoc = 0x78F4C3C0
	jmpOff = secondPatchLoc - firstPatchLoc
	patches = {
		firstPatchLoc: nasm("jmp %s" % hex(jmpOff)),
		secondPatchLoc: nasm('''
			mov r14, qword [rcx + 8]         ; r14 = ValueName->Buffer
			movzx   r14d, word [r14]         ; r14d = *r14
			cmp     r14d, 53h 
			jne short .syscall               ; if (r14d != 'S') goto .syscall
				mov     r13d, 1              ; else r13d = 1 
				mov     r12, r8              ; r12 = SystemData
			.syscall:
				mov     r10, rcx             ; r10 = ValueName
				mov     eax, 11Fh
				syscall
			cmp     r13d, 1
			jne short .out                   ; if (r13d != 1) goto .out
				;mov     byte [r12d], 1   ; *r12 = 1
				nop
			.out:
				retn
		''')
	}
	
	dll = list(open(r"C:\Windows\System32\ntdll-orig.dll", "rb").read())
	for loc, patch in patches.items():
		dll = applyPatch(dll, loc, patch)
		
	open(r"E:\dev\cpp\ntdll\test_ntdll3.dll", "wb").write(bytes(dll))

if __name__ == "__main__":
	main()