#include <Security/RuntimeSecurity.hpp>
#include <Security/BretHart.hpp>

#include <UserExperience/UserInterface.hpp>
#include <Networking/TCPClient.hpp>

// Global accessor to security instance.
Security::RuntimeSecurityPtr Protection = std::make_unique<Security::RuntimeSecurity>();


namespace Security
{
	// Hooked functions.
	///////////////////////////////////////////////////////////

	decltype(&OpenProcess) oOpenProcess;
	HANDLE __stdcall Hooked_OpenProcess(DWORD AccessLevel, bool Inherit, DWORD ProcessId)
	{
		// Determine where the return address of the function actually points.
		void *Address                  = _ReturnAddress();
		MEMORY_BASIC_INFORMATION Query = Protection->QueryMemory(Address);

		// If the return address points outside of the loader module,
		// fail the function.
		HMODULE ReturnModule = (HMODULE)Query.AllocationBase;
		HMODULE LoaderModule = GetModuleHandleA(NULL);

		if(ReturnModule != LoaderModule)
		{
			Protection->SecurityCallback();

			[&](decltype(&OpenProcess) A)
			{
				// Again, let's meme anyone who tries to reverse this.
				uintptr_t NullPointer = *(uintptr_t *)0x00000000;
				A(NullPointer, NullPointer, NullPointer);
			}(OpenProcess);
		}

		// Call original function
		return oOpenProcess(AccessLevel, Inherit, ProcessId);
	}

	decltype(&ExitProcess) oExitProcess;
	void __stdcall Hooked_ExitProcess(DWORD ExitCode)
	{
		// Determine where the return address of the function actually points.
		void *Address                  = _ReturnAddress();
		MEMORY_BASIC_INFORMATION Query = Protection->QueryMemory(Address);

		// If the return address points outside of the loader module,
		// fail the function.
		HMODULE ReturnModule = (HMODULE)Query.AllocationBase;
		HMODULE LoaderModule = GetModuleHandleA(NULL);

		if(ReturnModule != LoaderModule) 
		{
			Protection->SecurityCallback();

			[&](decltype(&ExitProcess) A)
			{
				// Again, let's meme anyone who tries to reverse this.
				uintptr_t NullPointer = *(uintptr_t *)0x00000000;
				A(NullPointer);
			}(oExitProcess);
		}

		// Call original function
		oExitProcess(ExitCode);
	}

	decltype(&recv) oWSARecv;
	int __stdcall Hooked_WSARecv(SOCKET Socket, char *Buffer, int Length, int Flags)
	{
		// Determine where the return address of the function actually points.
		void *Address                  = _ReturnAddress();
		MEMORY_BASIC_INFORMATION Query = Protection->QueryMemory(Address);

		// If the return address points outside of the loader module,
		// fail the function.
		HMODULE ReturnModule = (HMODULE)Query.AllocationBase;
		HMODULE LoaderModule = GetModuleHandleA(NULL);

		if(ReturnModule != LoaderModule)
			// Let's meme anyone who tries to reverse this.
			return []() { Protection->SecurityCallback(); return -1; }();

		// Call original function
		return oWSARecv(Socket, Buffer, Length, Flags);
	}

	decltype(&send) oWSASend;
	int __stdcall Hooked_WSASend(SOCKET Socket, char *Buffer, int Length, int Flags)
	{
		// Determine where the return address of the function actually points.
		void *Address                  = _ReturnAddress();
		MEMORY_BASIC_INFORMATION Query = Protection->QueryMemory(Address);

		// If the return address points outside of the loader module,
		// fail the function.
		HMODULE ReturnModule = (HMODULE)Query.AllocationBase;
		HMODULE LoaderModule = GetModuleHandleA(NULL);

		if(ReturnModule != LoaderModule)
			// Let's meme anyone who tries to reverse this.
			return []() { Protection->SecurityCallback(); return -1; }();

		// Call original function
		return oWSASend(Socket, Buffer, Length, Flags);
	}

	// The following functions are only called internally.
	///////////////////////////////////////////////////////////

	// Sick macros, retard.
#define CreateMinHook()		MH_STATUS Status; Status = MH_Initialize();
#define CheckStatus()			if(Status != MH_OK) { return false; }
#define SafeCallTo(Function)	Status = Function; CheckStatus();

	bool RuntimeSecurity::ApplyApiHooks()
	{
		// Make sure that MinHook is initialized properly.
		CreateMinHook();
		CheckStatus();

		// Apply any hooks.
		SafeCallTo(MH_CreateHook(&OpenProcess, Hooked_OpenProcess, (void **)&oOpenProcess));
		SafeCallTo(MH_EnableHook(&OpenProcess));

		SafeCallTo(MH_CreateHook(&ExitProcess, Hooked_ExitProcess, (void **)&oExitProcess));
		SafeCallTo(MH_EnableHook(&ExitProcess));

		SafeCallTo(MH_CreateHook(&recv, Hooked_WSARecv, (void **)&oWSARecv));
		SafeCallTo(MH_EnableHook(&recv));

		SafeCallTo(MH_CreateHook(&send, Hooked_WSASend, (void **)&oWSASend));
		SafeCallTo(MH_EnableHook(&send));

		return true;
	}

	void RuntimeSecurity::PatchDebugFunctions()
	{
		WRAP_IF_RELEASE(VM_EAGLE_WHITE_START);
		WRAP_IF_RELEASE(STR_ENCRYPT_START);

		HMODULE Module = GetModuleHandleA("ntdll.dll");

		if(!Module)
			ERROR_ASSERT("[000F:00001A00] Failed to initialize. Please contact an administrator.");

		// Grab exports from ntdll.dll
		uintptr_t Export_DbgUiRemoteBreakin = (uintptr_t)GetProcAddress(Module, "DbgUiRemoteBreakin");
		uintptr_t Export_DbgBreakPoint      = (uintptr_t)GetProcAddress(Module, "DbgBreakPoint");
		
		// Most plugins for OllyDBG / IDA only fix DbgUiRemoteBreakin/DbgBreakPoint,
		// however, NtContinue is never touched although it is used.
		// This should prevent any such plugins from effectively attaching the debugger.
		// NOTE: This does not work on x64dbg for whatever reason..
		uintptr_t Export_NtContinue         = (uintptr_t)GetProcAddress(Module, "NtContinue");

		// Our small shellcode that calls oExitProcess.
		uint8_t Shellcode[] = {
			0xB8, 0x00, 0x00, 0x00, 0x00,	// mov eax, dword ptr[0]
			0xFF, 0xE0						// jmp eax
		};

		std::memcpy(&Shellcode[1], oExitProcess, sizeof uint32_t);

		// Ensure that the program gets closed if a debugger is attached.
		uintptr_t Exports[] = {
			Export_DbgUiRemoteBreakin,
			Export_DbgBreakPoint,
			//Export_NtContinue // This causes a lot of crashes ATM while debugging, leave this out till release.
		};

		for(auto &It : Exports)
		{
			DWORD OldProtection;
			if(!VirtualProtect((void *)It, sizeof Shellcode, PAGE_EXECUTE_READWRITE, &OldProtection))
				ERROR_ASSERT("[000F:00001A00] Failed to initialize. Please contact an administrator.");

			// Patch to __asm { jmp oExitProcess; };
			*(uint8_t *)It        = 0xE9;
			*(uint32_t *)(It + 1) = (uintptr_t)oExitProcess;

			VirtualProtect((void *)It, sizeof Shellcode, OldProtection, &OldProtection);
		}
		
		WRAP_IF_RELEASE(STR_ENCRYPT_END);
		WRAP_IF_RELEASE(VM_EAGLE_WHITE_END);
	}

	void RuntimeSecurity::DispatchSecurityThreads()
	{
		WRAP_IF_RELEASE(MUTATE_START);

		std::thread DebugThread	(&RuntimeSecurity::CheckForDebugger, this);			DebugThread.detach();
		std::thread VMThread	(&RuntimeSecurity::CheckForVirtualMachine, this);	VMThread.detach();
		std::thread DriverThread(&RuntimeSecurity::CheckForDrivers, this);			DriverThread.detach();
		std::thread TamperThread(&RuntimeSecurity::CheckForThreads, this);			TamperThread.detach();

		WRAP_IF_RELEASE(MUTATE_END);
	}

	// The following functions are only called internally.
	// The reason these are seperate is in order to maintain
	// code readability.
	///////////////////////////////////////////////////////////

	void RuntimeSecurity::CheckForVirtualMachine()
	{

	}

	void RuntimeSecurity::CheckForDebugger()
	{
		// Read the PEB from the TIB.
		// Offset for x86 is 0x30 ; mov ..., dword ptr fs:[0x30]
		// Offset for x64 is 0x60 ; mov ..., qword ptr gs:[0x60]

		//PEB *ProcessEnvBlock = (PEB *)__readgsqword(0x60);
		//
		//if(ProcessEnvBlock->BeingDebugged)
		//	SecurityCallback();
	}

	void RuntimeSecurity::CheckForDrivers()
	{

	}

	void RuntimeSecurity::CheckForThreads()
	{

	}

	// The following functions are exposed publicly.
	///////////////////////////////////////////////////////////

	bool RuntimeSecurity::Start()
	{
		// If hooking API functions fails, exit the program.
		if(!ApplyApiHooks())
			return false;

		// Dispatch threads before patching NtContinue & co.
		DispatchSecurityThreads();

		// Patch DbgUiRemoteBreakin, DbgBreakPoint, NtContinue
		// This also fucks up detours for some reason... only extra protection :-)
		PatchDebugFunctions();

		return true;
	}

	HardwareIdentifier RuntimeSecurity::GetHardwareId()
	{
		return HardwareIdentifier{};
	}

	MEMORY_BASIC_INFORMATION RuntimeSecurity::QueryMemory(void *Address)
	{
		MEMORY_BASIC_INFORMATION Result;

		VirtualQuery(Address, &Result, sizeof MEMORY_BASIC_INFORMATION);

		return Result;
	}

	__forceinline void RuntimeSecurity::SecurityCallback()
	{
		// TODO: Implement something that bans the user or notifies
		// you if someone tampers with the loader.
		// For convenience sake, I'll just put a small .WAV file of
		// Bret Hart screaming in here for now :^)
		PlaySoundA((char *)Resource::BretWAV, NULL, SND_MEMORY);
		oExitProcess(0);
	}
}