TotalA-Engine-Notes/TotalA-Dev-Projects/AI Lua Scripting/Patching.hpp

#pragma once

#include "Defines.hpp"
#include "BuiltIn.hpp"

#include <Windows.h>

enum PatchType
{
	PATCHTYPE_BYTEPATCH,
	PATCHTYPE_INTERCEPT
};

enum JumpType
{
	JUMPTYPE_JMP,
	JUMPTYPE_CALL
};

namespace PatchFix
{
	namespace Patching
	{
		struct Patch
		{
		public:
			VoidPtr ReturnAddr;
		private:
			BytePtr PatchBytes;
			BytePtr OriginalBytes;
			Dword PatchType;
			Dword JumpType;
			Dword ByteCount;
			VoidPtr PatchAddr;
			VoidPtr InterceptAddr;

			Dword LastProtect;

		public:
			Patch();
			void Setup(Dword PatchType, VoidPtr Bytes, Dword ByteCount, Dword JumpType, VoidPtr PatchAddress, VoidPtr InterceptAddress, VoidPtr ReturnAddress);
			void Apply();
			void Revert();

		private:
			bool ReProtect(VoidPtr Addr, Dword Length);
			bool UnProtect(VoidPtr Addr, Dword Length);
			void BackupOriginal(VoidPtr Addr, Dword Length);
		};
	}
}



PatchFix::Patching::Patch::Patch()
{
	PatchBytes = null;
	OriginalBytes = null;
	PatchType = null;
	JumpType = null;
	ByteCount = null;
	PatchAddr = null;
	InterceptAddr = null;
	ReturnAddr = null;
	LastProtect = null;
}

void PatchFix::Patching::Patch::Setup(Dword PatchType, VoidPtr Bytes, Dword ByteCount, Dword JumpType, VoidPtr PatchAddress, VoidPtr InterceptAddress, VoidPtr ReturnAddress)
{
	this->PatchType = PatchType;
	this->ByteCount = ByteCount;
	this->JumpType = JumpType;
	this->PatchAddr = PatchAddress;
	this->InterceptAddr = InterceptAddress;
	this->ReturnAddr = ReturnAddress;

	BytePtr BytesPtr;
	DwordPtr PatchBytesJmpPtr;
	int JumpDistance;

	if (PatchType == PATCHTYPE_BYTEPATCH)
	{
		BackupOriginal(PatchAddress, ByteCount);

		PatchBytes = (BytePtr)malloc(ByteCount);

		if (PatchBytes != null)
		{
			BytesPtr = (BytePtr)Bytes;

			for (Dword I = 0; I < ByteCount; I++)
			{
				PatchBytes[I] = BytesPtr[I];
			}
		}
	}
	else if (PatchType == PATCHTYPE_INTERCEPT)
	{
		this->ByteCount = 5;

		BackupOriginal(PatchAddress, this->ByteCount);

		PatchBytes = (BytePtr)malloc(5);

		if (PatchBytes != null)
		{
			if (JumpType == JUMPTYPE_JMP)
			{
				PatchBytes[0] = 0xE9;
			}
			else if (JumpType == JUMPTYPE_CALL)
			{
				PatchBytes[0] = 0xE8;
			}

			PatchBytesJmpPtr = (DwordPtr)&PatchBytes[1];

			JumpDistance = (int)InterceptAddress - (int)PatchAddress - 5;

			*PatchBytesJmpPtr = JumpDistance;
		}
	}
}

void PatchFix::Patching::Patch::Apply()
{
	BytePtr PatchBytesAddrPtr = (BytePtr)PatchAddr;

	if (UnProtect(PatchAddr, ByteCount))
	{
		for (Dword I = 0; I < ByteCount; I++)
		{
			PatchBytesAddrPtr[I] = PatchBytes[I];
		}

		ReProtect(PatchAddr, ByteCount);
	}
}

void PatchFix::Patching::Patch::Revert()
{
	BytePtr PatchBytesAddrPtr = (BytePtr)PatchAddr;

	if (UnProtect(PatchAddr, ByteCount))
	{
		for (Dword I = 0; I < ByteCount; I++)
		{
			PatchBytesAddrPtr[I] = OriginalBytes[I];
		}

		ReProtect(PatchAddr, ByteCount);
	}
}

bool PatchFix::Patching::Patch::ReProtect(VoidPtr Addr, Dword Length)
{
	return PatchFix::BuiltIn::TA_VirtualProtect(Addr, Length, LastProtect, &LastProtect);
}

bool PatchFix::Patching::Patch::UnProtect(VoidPtr Addr, Dword Length)
{
	return PatchFix::BuiltIn::TA_VirtualProtect(Addr, Length, PAGE_EXECUTE_READWRITE, &LastProtect);
}

void PatchFix::Patching::Patch::BackupOriginal(VoidPtr Addr, Dword Length)
{
	BytePtr OriginalSource = (BytePtr)Addr;

	OriginalBytes = (BytePtr)malloc(Length);

	if (OriginalBytes != null)
	{
		for (Dword I = 0; I < Length; I++)
		{
			OriginalBytes[I] = OriginalSource[I];
		}
	}
}