#include <Windows.h>
#include <tlhelp32.h>
#include <string>
#include "MinHook.h"
#include "scanner.h"
#include "nt.h"

#define FUNCTION_SIG "48 8b 91 ? ? ? ? 48 85 d2 74 ? 48 8b 51"

#define FUNCTION_COUNT 9

def_LdrRegisterDllNotification Org_LdrRegisterDllNotification = NULL;
def_LdrUnregisterDllNotification Org_LdrUnregisterDllNotification = NULL;
def_FreeLibrary Org_FreeLibrary = NULL;
inline PVOID Cookie = NULL;
inline void* sigpointer =  nullptr;
extern "C" uintptr_t functions[FUNCTION_COUNT] = { 0 };
std::string functionNames[] =
{
	"NvFBC_Create",
	"NvFBC_CreateEx",
	"NvFBC_Enable",
	"NvFBC_GetSDKVersion",
	"NvFBC_GetStatus",
	"NvFBC_GetStatusEx",
	"NvFBC_SetGlobalFlags",
	"NvFBC_UnloadNvAPI",
	"NvOptimusEnablement"
};


typedef enum _NVFBCRESULT
{
	NVFBC_SUCCESS = 0,
	NVFBC_ERROR_GENERIC = -1,                     /**< Unexpected failure in NVFBC. */
	NVFBC_ERROR_INVALID_PARAM = -2,               /**< One or more of the paramteres passed to NvFBC are invalid [This include NULL pointers]. */
	NVFBC_ERROR_INVALIDATED_SESSION = -3,         /**< NvFBC session is invalid. Client needs to recreate session. */
	NVFBC_ERROR_PROTECTED_CONTENT = -4,           /**< Protected content detected. Capture failed. */
	NVFBC_ERROR_DRIVER_FAILURE = -5,              /**< GPU driver returned failure to process NvFBC command. */
	NVFBC_ERROR_CUDA_FAILURE = -6,              /**< CUDA driver returned failure to process NvFBC command. */
	NVFBC_ERROR_UNSUPPORTED = -7,              /**< API Unsupported on this version of NvFBC. */
	NVFBC_ERROR_HW_ENC_FAILURE = -8,             /**< HW Encoder returned failure to process NVFBC command. */
	NVFBC_ERROR_INCOMPATIBLE_DRIVER = -9,         /**< NVFBC is not compatible with this version of the GPU driver. */
	NVFBC_ERROR_UNSUPPORTED_PLATFORM = -10,       /**< NVFBC is not supported on this platform. */
	NVFBC_ERROR_OUT_OF_MEMORY = -11,             /**< Failed to allocate memory. */
	NVFBC_ERROR_INVALID_PTR = -12,             /**< A NULL pointer was passed. */
	NVFBC_ERROR_INCOMPATIBLE_VERSION = -13,       /**< An API was called with a parameter struct that has an incompatible version. Check dwVersion field of paramter struct. */
	NVFBC_ERROR_OPT_CAPTURE_FAILURE = -14,        /**< Desktop Capture failed. */
	NVFBC_ERROR_INSUFFICIENT_PRIVILEGES = -15,   /**< User doesn't have appropriate previlages. */
	NVFBC_ERROR_INVALID_CALL = -16,               /**< NVFBC APIs called in wrong sequence. */
	NVFBC_ERROR_SYSTEM_ERROR = -17,               /**< Win32 error. */
	NVFBC_ERROR_INVALID_TARGET = -18,             /**< The target adapter idx can not be used for NVFBC capture. It may not correspond to an NVIDIA GPU, or may not be attached to desktop. */
	NVFBC_ERROR_DYNAMIC_DISABLE = -20,            /**< NvFBC is dynamically disabled. Cannot continue to capture */
} NVFBCRESULT;

// [index: 000] [ordinal: 001] NvFBC_Create
#pragma comment(linker, "/EXPORT:NvFBC_Create=func_9be0fbdc869f417b9f29c27e9d9c0cfd,@1")
extern "C" void func_9be0fbdc869f417b9f29c27e9d9c0cfd();

// [index: 001] [ordinal: 002] NvFBC_CreateEx
#pragma comment(linker, "/EXPORT:NvFBC_CreateEx=func_92a7155ee30d412ea0e85c05d5ef3cd8,@2")
extern "C" void func_92a7155ee30d412ea0e85c05d5ef3cd8();

// [index: 002] [ordinal: 003] NvFBC_Enable
#pragma comment(linker, "/EXPORT:NvFBC_Enable=func_358080a229af49bb96c3968f7c8d7444,@3")
extern "C" void func_358080a229af49bb96c3968f7c8d7444();

// [index: 003] [ordinal: 004] NvFBC_GetSDKVersion
#pragma comment(linker, "/EXPORT:NvFBC_GetSDKVersion=func_a3384a768753452cbccab384f5f330c1,@4")
extern "C" void func_a3384a768753452cbccab384f5f330c1();

// [index: 004] [ordinal: 005] NvFBC_GetStatus
#pragma comment(linker, "/EXPORT:NvFBC_GetStatus=func_c5a74b9ae49d4851930995268cee4e2f,@5")
extern "C" void func_c5a74b9ae49d4851930995268cee4e2f();

// [index: 005] [ordinal: 006] NvFBC_GetStatusEx
#pragma comment(linker, "/EXPORT:NvFBC_GetStatusEx=func_7a0e96ba3cf34b0e9815e3a16dc3d347,@6")
extern "C" void func_7a0e96ba3cf34b0e9815e3a16dc3d347();

// [index: 006] [ordinal: 007] NvFBC_SetGlobalFlags
#pragma comment(linker, "/EXPORT:NvFBC_SetGlobalFlags=func_2fd4dec9416b42de88a5afbe23ed4a8d,@7")
extern "C" void func_2fd4dec9416b42de88a5afbe23ed4a8d();

// [index: 007] [ordinal: 008] NvFBC_UnloadNvAPI
#pragma comment(linker, "/EXPORT:NvFBC_UnloadNvAPI=func_e8c80ebd98bb45e0b591f52e5fbfa998,@8")
extern "C" NVFBCRESULT func_e8c80ebd98bb45e0b591f52e5fbfa998() {
	return NVFBC_ERROR_GENERIC; //malware reject
};

// [index: 008] [ordinal: 009] NvOptimusEnablement
#pragma comment(linker, "/EXPORT:NvOptimusEnablement=func_5e40589f71e44de1827491903b6e99af,@9")
extern "C" void func_5e40589f71e44de1827491903b6e99af();


void* Org_GetWindowDisplayAffinity = nullptr;
BOOL WINAPI Hook_GetWindowDisplayAffinity(IN HWND hWnd, OUT DWORD* pwdAffinity) {
	*pwdAffinity = WDA_NONE;
	return TRUE;
}

void* Org_Module32FirstW = nullptr;
BOOL WINAPI Hook_Module32FirstW(IN HANDLE hSnapshot, IN OUT LPMODULEENTRY32W lpme) {
	return FALSE;
}

BOOL Hook_FreeLibrary(IN HMODULE hLibModule) {

	if (hLibModule == GetModuleHandleA("NvFBC64.dll"))
	{
		//MessageBoxA(nullptr,"SELF FREELIBRARY","WARN",MB_OK);
		//Org_LdrUnregisterDllNotification(Cookie);
		//MH_DisableHook(MH_ALL_HOOKS);
		//MH_Uninitialize();
		return false; //reject // iam malware
		//Sleep(1000);
	}
	return Org_FreeLibrary(hLibModule);
}

void* Org_sub_7FFE7411BB10 = nullptr;
__int64 __fastcall Hook_sub_7FFE7411BB10(__int64 a1)
{
	/*
	return *(unsigned __int64*)(a1 + 10896)
		&& 
		(*(unsigned int*)(*(unsigned __int64*)(a1 + 24) + 368LL) || *(unsigned int*)(*(unsigned __int64*)(a1 + 24) + 640LL));*/
	*(unsigned int*)(*(unsigned __int64*)(a1 + 24) + 368LL) = 0;
	*(unsigned int*)(*(unsigned __int64*)(a1 + 24) + 640LL) = 0;
	return 0;
}

void CALLBACK DLLNotification(ULONG Reason, PLDR_DLL_NOTIFICATION_DATA NotificationData, PVOID Context) {
	switch (Reason)
	{
		case LDR_DLL_NOTIFICATION_REASON_LOADED:
		{
			if (wcsstr(NotificationData->Loaded.FullDllName->Buffer, L"nvd3dumx.dll") != NULL) {
				sigpointer = sig(GetModuleHandleA("nvd3dumx.dll"), FUNCTION_SIG);
				if (sigpointer == nullptr)
				{
					MessageBoxA(nullptr, "nvd3dumx sig not found!", "ERROR", MB_ICONERROR | MB_OK);
				}
				else
				{
					MH_CreateHook(sigpointer, &Hook_sub_7FFE7411BB10, &Org_sub_7FFE7411BB10);
				}
				MH_EnableHook(sigpointer);
			}
			break;
		}
		case LDR_DLL_NOTIFICATION_REASON_UNLOADED:
		{
			
			if (wcsstr(NotificationData->Unloaded.FullDllName->Buffer, L"nvd3dumx.dll") != NULL)
			{
				/*
				Org_LdrUnregisterDllNotification(Cookie);
				MH_DisableHook(MH_ALL_HOOKS);
				MH_Uninitialize();
				*/
				if (sigpointer!=nullptr)
				{
					MH_DisableHook(sigpointer);
				}
			}
			
			break;
		}
	}
	/*
	if (Reason == LDR_DLL_NOTIFICATION_REASON_LOADED) {
		//wprintf(L"[LdrDllNotification] %s\n", NotificationData->Loaded.FullDllName->Buffer);
		
	}
	//return;
	*/
}

BOOL WINAPI DllMain(HINSTANCE hinstModule, DWORD dwReason, LPVOID lpvReserved)
{
	if (dwReason == DLL_PROCESS_ATTACH)
	{
		DisableThreadLibraryCalls(hinstModule);
		auto moduleHandle = LoadLibraryA("NvFBC64_original.dll");
		if (moduleHandle != nullptr)
		{
			for (auto i = 0; i < FUNCTION_COUNT; i++)
			{
				auto address = GetProcAddress(moduleHandle, functionNames[i].c_str());
				functions[i] = reinterpret_cast<uintptr_t>(address);
			}
		}
		
		if (MH_Initialize() != MH_OK) {
			MessageBoxA(nullptr, "MH Init Error!", "ERROR", MB_ICONERROR | MB_OK);
			exit(1);
		}


		MH_CreateHookApi(L"user32.dll","GetWindowDisplayAffinity",&Hook_GetWindowDisplayAffinity,&Org_GetWindowDisplayAffinity);
		MH_CreateHookApi(L"kernel32.dll", "Module32FirstW", &Hook_Module32FirstW, &Org_Module32FirstW);
		MH_CreateHookApi(L"kernel32.dll", "FreeLibrary", &Hook_FreeLibrary, reinterpret_cast<LPVOID*>(&Org_FreeLibrary));
		MH_EnableHook(MH_ALL_HOOKS);

		//MessageBoxA(nullptr, "DEBUG", "HOOK RD", MB_OK);

		Org_LdrRegisterDllNotification = (def_LdrRegisterDllNotification)GetProcAddress(GetModuleHandleA("ntdll.dll"), "LdrRegisterDllNotification");
		Org_LdrUnregisterDllNotification = (def_LdrUnregisterDllNotification)GetProcAddress(GetModuleHandleA("ntdll.dll"), "LdrUnregisterDllNotification");
		Org_LdrRegisterDllNotification(0, DLLNotification, NULL, &Cookie);
		//LoadLibraryA("nvd3dumx.dll");
		/*
		void* SigPointer = sig(GetModuleHandleA("nvd3dumx.dll"),FUNCTION_SIG);
		if (SigPointer == nullptr)
		{
			MessageBoxA(nullptr, "nvd3dumx sig not found!", "ERROR",MB_ICONERROR|MB_OK);
		}
		else
		{
			MH_CreateHook(SigPointer,&Hook_sub_7FFE7411BB10,&Org_sub_7FFE7411BB10);
		}
		*/
		

		/*
		HANDLE NvModule = GetModuleHandleA("nvd3dumx.dll");
		if (NvModule == INVALID_HANDLE_VALUE)
		{
			MessageBoxA(nullptr, "NOMODULE","NOMODULE",MB_OK);
		}
		*/
	}
	else if (dwReason == DLL_PROCESS_DETACH)
	{
		Org_LdrUnregisterDllNotification(&Cookie);
		MH_DisableHook(MH_ALL_HOOKS);
		MH_Uninitialize();
	}

	return TRUE;
}