#include <Security/RuntimeSecurity.hpp>
#include <Security/SyscallManager.hpp>
#include <Security/FnvHash.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.
	///////////////////////////////////////////////////////////

#pragma optimize("", off)

	decltype(&OpenProcess) oOpenProcess;
	HANDLE __stdcall Hooked_OpenProcess(DWORD AccessLevel, bool Inherit, DWORD ProcessId)
	{
		WRAP_IF_RELEASE(VM_EAGLE_WHITE_START);

		// 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)
		{
			WRAP_IF_RELEASE(STR_ENCRYPT_START);
			Protection->SecurityCallback(__FUNCSIG__);
			WRAP_IF_RELEASE(STR_ENCRYPT_END);

			[&](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);

		WRAP_IF_RELEASE(VM_EAGLE_WHITE_END);
	}

	decltype(&ExitProcess) oExitProcess;
	void __stdcall Hooked_ExitProcess(DWORD ExitCode)
	{
		WRAP_IF_RELEASE(VM_EAGLE_WHITE_START);

		WRAP_IF_DEBUG(oExitProcess(ExitCode));

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

		WRAP_IF_RELEASE(VM_EAGLE_WHITE_END);
	}

	decltype(&recv) oWSARecv;
	int __stdcall Hooked_WSARecv(SOCKET Socket, char *Buffer, int Length, int Flags)
	{
		WRAP_IF_RELEASE(VM_EAGLE_WHITE_START);

		// 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);
		
		// Let's meme anyone who tries to reverse this.
		if(ReturnModule != LoaderModule)
		{		
			WRAP_IF_RELEASE(STR_ENCRYPT_START);
			return []() { Protection->SecurityCallback(__FUNCSIG__); return -1; }();
			WRAP_IF_RELEASE(STR_ENCRYPT_END);
		}

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

		WRAP_IF_RELEASE(VM_EAGLE_WHITE_END);
	}

	decltype(&send) oWSASend;
	int __stdcall Hooked_WSASend(SOCKET Socket, char *Buffer, int Length, int Flags)
	{
		WRAP_IF_RELEASE(VM_EAGLE_WHITE_START);

		// 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);

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

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

		WRAP_IF_RELEASE(VM_EAGLE_WHITE_END);
	}

#pragma optimize("", on)

	// 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();

#pragma optimize("", off)
	
	bool RuntimeSecurity::ApplyApiHooks()
	{
		WRAP_IF_RELEASE(MUTATE_START);

		// 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;

		WRAP_IF_RELEASE(MUTATE_END);
	}

#pragma optimize("", on)

	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");
		
		// 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 uintptr_t + 1, 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 uintptr_t + 1, 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::CheckForTampering,		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.
	///////////////////////////////////////////////////////////

#pragma optimize("", off)

	void RuntimeSecurity::CheckForVirtualMachine()
	{
		WRAP_IF_RELEASE(VM_EAGLE_BLACK_START);

		for(;;)
		{
			// Yeah, um, your code did absolutely fuck all in my analysis VM.
			int32_t VirtualMachineChecksum = 0x4000;

			WRAP_IF_RELEASE(
				CHECK_VIRTUAL_PC(VirtualMachineChecksum, 0x2000);

				WRAP_IF_RELEASE(STR_ENCRYPT_START);
				if(VirtualMachineChecksum != 0x2000)
					SecurityCallback(__FUNCSIG__);
				WRAP_IF_RELEASE(STR_ENCRYPT_END);
			);

			// Don't put too much stress on the CPU.
			Sleep(VirtualMachineChecksum);
		}


		WRAP_IF_RELEASE(VM_EAGLE_BLACK_END);
	}

	void RuntimeSecurity::CheckForDebugger()
	{
		WRAP_IF_RELEASE(VM_EAGLE_BLACK_START);
		WRAP_IF_RELEASE(STR_ENCRYPT_START);

		for(;;)
		{
			// 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(__FUNCSIG__);

			// TODO: Check for x64dbg window?
			/*
				------------------------------
				HWND: 000305A4
					HWND->m_Class = "ID"
					HWND->m_Text  = "Immunity
				------------------------------
				HWND: 00060574
					HWND->m_Class = || NON CONSTANT ||
					HWND->m_Text  = "x64dbg"
				------------------------------
				HWND: 002C0680
					HWND->m_Class = || NON CONSTANT ||
					HWND->m_Text  = "Progress Telerik Fiddler Web Debugger"
				------------------------------
				HWND: 000406E4
					HWND->m_Class = "OLLYDBG"
					HWND->m_Text  = "OllyDbg"
				------------------------------
			*/
			using WindowParams = std::pair<const char *, const char *>;
			static std::vector<WindowParams> BlackListedWindows = {
				{"ID",				"Immunity"},								// Immunity Debugger
				{"Qt5QWindowIcon",	"x64dbg"},									// x64dbg
				{"Qt5QWindowIcon",	"x32dbg"},									// x32dbg
				{"OLLYDBG",			"OllyDbg"},									// OllyDbg
				{nullptr,			"Progress Telerik Fiddler Web Debugger"},	// Telerik Fiddler
			};

			for(auto &It : BlackListedWindows)
			{
				// size_t Index = std::distance(...);

				if(FindWindowA(It.first, It.second))
					SecurityCallback(__FUNCSIG__);
			}

			// Don't put too much stress on the CPU.
			Sleep(150);
		}

		WRAP_IF_RELEASE(STR_ENCRYPT_END);
		WRAP_IF_RELEASE(VM_EAGLE_BLACK_END);
	}

	void RuntimeSecurity::CheckForDrivers()
	{
		WRAP_IF_RELEASE(VM_EAGLE_BLACK_START);
		
		// TODO: Check if test-signing mode is on
		// TODO: Check if safe-mode is on

		// TODO: Check for disallowed drivers
		for(;;)
		{

			// Don't put too much stress on the CPU.
			Sleep(150);
		}

		WRAP_IF_RELEASE(VM_EAGLE_BLACK_END);
	}

	void RuntimeSecurity::CheckForTampering()
	{
		WRAP_IF_RELEASE(VM_EAGLE_BLACK_START);

		for(;;)
		{
			int32_t CodeIntegrityChecksum = 0x2000;

			WRAP_IF_RELEASE(
				CHECK_CODE_INTEGRITY(CodeIntegrityChecksum, 0x4000);

				WRAP_IF_RELEASE(STR_ENCRYPT_START);
				if(CodeIntegrityChecksum != 0x4000)
					SecurityCallback(__FUNCSIG__);
				WRAP_IF_RELEASE(STR_ENCRYPT_END);
			);

			WRAP_IF_RELEASE(
				CHECK_PROTECTION(CodeIntegrityChecksum, 0x4000);

				WRAP_IF_RELEASE(STR_ENCRYPT_START);
				if(CodeIntegrityChecksum != 0x4000)
					SecurityCallback(__FUNCSIG__);
				WRAP_IF_RELEASE(STR_ENCRYPT_END);
			);

			// Don't put too much stress on the CPU.
			Sleep(CodeIntegrityChecksum);
		}


		WRAP_IF_RELEASE(VM_EAGLE_BLACK_END);
	}

#pragma optimize("", on)

	// 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{};
	}

#pragma optimize("", off)

	MEMORY_BASIC_INFORMATION RuntimeSecurity::QueryMemory(void *Address)
	{
		static auto ZwQueryVirtualMemory = Syscalls->Find<long(__stdcall *)(HANDLE, void *, int, void *, uint32_t, uint32_t *)>(FNV("ZwQueryVirtualMemory"));

		MEMORY_BASIC_INFORMATION Result{};
		NTSTATUS Status = ZwQueryVirtualMemory((HANDLE)-1, Address, 0, &Result, sizeof Result, nullptr);

		if(NT_ERROR(Status))
		{
			char ReasonParameter[64];

			WRAP_IF_DEBUG(sprintf_s(ReasonParameter, "[QueryMemory] NTSTATUS: %08x", Status));
			WRAP_IF_RELEASE(
				sprintf_s(ReasonParameter, "[00DF:%08x] There was an error with accessing a process.", Status);
				ERROR_ASSERT(ReasonParameter);
			);
			
			// yeet
			SecurityCallback(ReasonParameter);
		}

		return Result;
	}

	void RuntimeSecurity::SecurityCallback(const char *Reason)
	{
		WRAP_IF_RELEASE(VM_FISH_WHITE_START);

		static bool TriggeredCallback = false;

		if(!TriggeredCallback) 
		{
			WRAP_IF_RELEASE(STR_ENCRYPT_START);

			// You can use the reason parameters to debug the security in case
			// something weird starts going on with it.
			char Buffer[2048];
			sprintf_s(Buffer, "Security callback was called. Reason parameter: \"%s\"\n", Reason);

			MessageBoxA(0, Buffer, "Information", MB_ICONINFORMATION | MB_OK);


			// The process will straight up die on Release mode.
			// Compile with FuckMSVC to debug this.
			WRAP_IF_RELEASE(
				ExitProcess(rand() % RAND_MAX);
			);

			TriggeredCallback = true;
			WRAP_IF_RELEASE(STR_ENCRYPT_END);
		}

		WRAP_IF_RELEASE(VM_FISH_WHITE_END);
	}

#pragma optimize("", on)
}