/*

I don't know what to put here! :-P

Use code however you want, with common sense:
 - Give me 90% of any profit
 - Give credit
 - etc. :-)

*/


#include <Windows.h>
#include <Shlwapi.h>

#pragma comment(lib, "Kernel32.lib")

// For GetSystemMetrics/SAFE_MODE_CHECK
#pragma comment(lib, "User32.lib")

/*
 For Path* functions/wnsprintf with logging
 Needs to be linked with delayimp.lib and /DELAYLOAD:Shlwapi.dll
 Otherwise AppInit DLL won't be loaded in some cases: On XP, noticed
 with Pale Moon (only with SRP DLL checking!?) and QuickTime
*/
#pragma comment(lib, "Shlwapi.lib")


#ifndef PROCESS_DEP_ENABLE
#define PROCESS_DEP_ENABLE 1
#define PROCESS_DEP_DISABLE_ATL_THUNK_EMULATION 2
#endif

#define LOG_WAIT_TIME 256

#define FILE_SHARE_ALL (FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE)


typedef BOOL (WINAPI *GetProcDEPPolicy_f)(HANDLE, LPDWORD, PBOOL);
typedef BOOL (WINAPI *SetProcDEPPolicy_f)(DWORD);
typedef DEP_SYSTEM_POLICY_TYPE (WINAPI *GetSysDEPPolicy_f)(void);


// Globals only used with logging
HINSTANCE thisDLL;

GetProcDEPPolicy_f GetProcessDEP;
SetProcDEPPolicy_f SetProcessDEP;


#if defined(LOGGING) || defined(NO_UNLOAD)
#define UNLOAD() 0
#elif defined(SANDBOXIE)
#define UNLOAD() !GetModuleHandleW(L"SbieDll.dll")
#else
#define UNLOAD() 1
#endif

#ifdef SAFE_MODE_CHECK
#define IN_SAFE_MODE() GetSystemMetrics(SM_CLEANBOOT)
#else
#define IN_SAFE_MODE() 0
#endif


#ifdef NO_NX_CHECK
// NX_COMPAT already handled on Vista SP1 and later, so no point in checking
#define Is_NX_COMPAT() 0
#else
#define RVA2PTR(base, rva) (((PBYTE) base) + rva)

__forceinline BOOL Is_NX_COMPAT(void)
{
	// GetModuleHandle's pointer to the loaded executable is just like reading the file's contents to check the header
	PIMAGE_DOS_HEADER DOS_header = (PIMAGE_DOS_HEADER) GetModuleHandleW(NULL);
	PIMAGE_NT_HEADERS NT_header  = (PIMAGE_NT_HEADERS) RVA2PTR(DOS_header, DOS_header->e_lfanew);

	if (NT_header->Signature == IMAGE_NT_SIGNATURE && NT_header->OptionalHeader.DllCharacteristics & IMAGE_DLLCHARACTERISTICS_NX_COMPAT)
	{
		return TRUE;
	}

	return FALSE;
}
#endif


#ifdef OPTOUT_CHECK_FOR_WIN8
__forceinline BOOL OptOut_SysDEP(HANDLE hKernel32)
{
	GetSysDEPPolicy_f GetSystemDEP = (GetSysDEPPolicy_f) GetProcAddress(hKernel32, "GetSystemDEPPolicy");

	return (GetSystemDEP() == DEPPolicyOptOut);
}
#else
#define OptOut_SysDEP(UNUSED) 1
#endif


#ifdef LOGGING
DWORD WINAPI LogDEP(DWORD orig_flags)
{
	BOOL NX_on_XP = orig_flags & PROCESS_DEP_DISABLE_ATL_THUNK_EMULATION;

	DWORD flags;
	BOOL permanent;

	LPCSTR msg;

	HANDLE log_file;
	char logDLL_path[MAX_PATH + sizeof("PermanentDEP.txt")];


	GetProcessDEP(GetCurrentProcess(), &flags, &permanent);

	if (permanent == TRUE && flags == PROCESS_DEP_ENABLE)
	{
#ifdef FULL_EMET_CHECK
		char EMET_Settings[256];

		if (GetEnvironmentVariableA("EMET_Settings", EMET_Settings, 256) && StrStrA(EMET_Settings, "DEP:1"))
#else
		if (GetModuleHandleW(L"EMET.dll"))
#endif
		{
			msg = NX_on_XP ? "*NX,EMET" : "-NX,EMET";

			goto Log;
		}
	}

	Sleep(LOG_WAIT_TIME);

	GetProcessDEP(GetCurrentProcess(), &flags, &permanent);

	if (permanent != TRUE)
	{
		// flags == 0, at this point, SHOULD only happen on XP with OptIn
		// Then we'd only be here with orig ATL_THUNK flag because of NX_COMPAT
		if (flags || NX_on_XP)
		{
			msg = NX_on_XP ? "+NX (XP)" : "<No NX>";

			SetProcessDEP(orig_flags);
		}
		else
		{
			msg = "**Off**";
		}
	}
	else
	{
		if (flags & PROCESS_DEP_DISABLE_ATL_THUNK_EMULATION)
		{
			msg = "Set-ATL";
		}
		else if (flags)
		{
			msg = NX_on_XP ? "*NX,Set" : "-NX,Set";
		}
		else
		{
			msg = "**OFF!**";
		}
	}
Log:
	GetModuleFileNameA(thisDLL, logDLL_path, MAX_PATH);
	PathRemoveFileSpecA(logDLL_path);
	PathAppendA(logDLL_path, "PermanentDEP.txt");

	if (log_file = CreateFileA(logDLL_path, FILE_APPEND_DATA | FILE_READ_DATA, FILE_SHARE_ALL, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL))
	{
		char exe_name[MAX_PATH + 1];

		char log_entry[MAX_PATH + 16];
		int entry_len;

		GetModuleFileNameA(NULL, exe_name, MAX_PATH);
		PathStripPathA(exe_name);

		entry_len = wnsprintfA(log_entry, MAX_PATH + 15, "%-9s %s\r\n", msg, exe_name);

		WriteFile(log_file, log_entry, entry_len, &entry_len, NULL);
		CloseHandle(log_file);
	}

	return 0;
}
#endif


__forceinline void CheckAndSetPermanentDEP(void)
{
	HANDLE hKernel32 = GetModuleHandleW(L"kernel32.dll");

	DWORD flags;
	BOOL permanent;

#ifndef LOGGING
	GetProcDEPPolicy_f GetProcessDEP;
	SetProcDEPPolicy_f SetProcessDEP;
#endif

	if (hKernel32)
	{
		GetProcessDEP = (GetProcDEPPolicy_f) GetProcAddress(hKernel32, "GetProcessDEPPolicy");

		if (GetProcessDEP)
		{
			GetProcessDEP(GetCurrentProcess(), &flags, &permanent);

			if (permanent != TRUE && !IN_SAFE_MODE())
			{
				if (Is_NX_COMPAT())
				{
					flags = PROCESS_DEP_ENABLE | PROCESS_DEP_DISABLE_ATL_THUNK_EMULATION;
				}
				else if (!(flags & PROCESS_DEP_ENABLE && (!permanent || OptOut_SysDEP(hKernel32))))
				{
					return;
				}


				SetProcessDEP = (SetProcDEPPolicy_f) GetProcAddress(hKernel32, "SetProcessDEPPolicy");

				if (SetProcessDEP)
				{
#ifdef LOGGING
					CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE) LogDEP, (LPVOID) flags, 0, NULL);
#else
					SetProcessDEP(flags);
#endif
				}
			}
		}
	}
}


BOOL WINAPI DllMain(HINSTANCE hiDLL, DWORD dwReason, LPVOID lpReserved)
{
	if (dwReason == DLL_PROCESS_ATTACH)
	{
#ifdef LOGGING
		thisDLL = hiDLL;
#endif
		CheckAndSetPermanentDEP();

		if (UNLOAD())
		{
			return FALSE;
		}
	}

	return TRUE;
}