/*

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 <Psapi.h>

#pragma comment(lib, "User32.lib")
#pragma comment(lib, "Advapi32.lib")
#pragma comment(lib, "Psapi.lib")


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

#define PROCESS_FLAGS (PROCESS_QUERY_INFORMATION | PROCESS_VM_READ | PROCESS_VM_OPERATION | PROCESS_CREATE_THREAD)


typedef DEP_SYSTEM_POLICY_TYPE (WINAPI *GetSysDEPPolicy_f)(void);


const char usage_txt[] =
	"SetDEP v1.0.0, by Matt Wilmas\n"
	" Set DEP options for a process.\n\n"
	"Usage: SetDEP [-c] [-d | -a] <process | PID>  [| more]\n"
	" Enables DEP (permanently) by default (without -d or -a)\n\n"
	" -d  Disable DEP if it's not \"permanently\" enabled\n"
	" -a  Disable \"ATL thunk emulation\" (like /NXCOMPAT, or AlwaysOn System DEP)\n\n"
	" -c  Console output of usage text; command status always goes to console\n"
	"      You won't see output in a Command Prompt unless you pipe to  | more\n"
	"      (see or use cliDEP.bat)";


LPCSTR SysDEP_names[] = { "AlwaysOff", "AlwaysOn", "OptIn", "OptOut" };


#define printf(...) my_printf(__VA_ARGS__)

int my_printf(LPCSTR format, ...)
{
	char buf[1024];
	int len;

	va_list args;

	va_start(args, format);
	len = wvsprintfA(buf, format, args);
	va_end(args);

	WriteFile(GetStdHandle(STD_OUTPUT_HANDLE), buf, len, &len, NULL);

	return len;
}

BOOL EnableDebugPrivilege(void)
{
	HANDLE hToken;
	TOKEN_PRIVILEGES tkp; 

	if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
	{
		if (LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &tkp.Privileges[0].Luid))
		{
			tkp.PrivilegeCount = 1;
			tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

        		if (AdjustTokenPrivileges(hToken, FALSE, &tkp, 0, NULL, NULL) && GetLastError() == ERROR_SUCCESS)
			{        
				CloseHandle(hToken);

				return TRUE;
			}
		}

		CloseHandle(hToken);
	}

	return FALSE;
}

void Message(LPCSTR message, BOOL console)
{
	if (console)
	{
		printf("%s", message);
	}
	else
	{
		MessageBoxA(NULL, message, "SetDEP", MB_OK);
	}
}

void Usage(BOOL console)
{
	HANDLE hKernel32 = GetModuleHandleW(L"kernel32.dll");
	GetSysDEPPolicy_f GetSystemDEP = (GetSysDEPPolicy_f) GetProcAddress(hKernel32, "GetSystemDEPPolicy");

	char buf[128], output[1024];
	LPSTR SysDEP_msg = buf;

	if (GetSystemDEP)
	{
		DEP_SYSTEM_POLICY_TYPE SysDEP_type = GetSystemDEP();
		LPCSTR type_msg;

		OSVERSIONINFO versionInfo;

		switch (SysDEP_type)
		{
			case DEPPolicyAlwaysOff:
			case DEPPolicyAlwaysOn:
				type_msg = "and SetDEP will have no effect";
				break;
			case DEPPolicyOptIn:
				versionInfo.dwOSVersionInfoSize = sizeof(versionInfo);

				GetVersionEx(&versionInfo);

				if (versionInfo.dwMajorVersion < 6)
				{
					type_msg = "is useful, on XP";
				}
				else
				{
					type_msg = "is little to no use, on Vista and later";
				}

				break;
			case DEPPolicyOptOut:
				type_msg = "is useful";
				break;
			default:
				return; // Shouldn't happen, prevent crash
		}

		wsprintfA(SysDEP_msg, "System DEP: %s (Permanent DEP DLL %s)", SysDEP_names[SysDEP_type], type_msg);
	}
	else
	{
		SysDEP_msg = "Note: This system doesn't have the DEP functions!";
	}

	wsprintfA(output, "%s\n\n\n%s", usage_txt, SysDEP_msg);
	Message(output, console);
}

HANDLE OpenProcessByName(LPCSTR target, LPDWORD pid)
{
	DWORD PIDs[1024];
	DWORD size;

	UINT target_len = lstrlenA(target);

	UINT i;

	if (!EnumProcesses(PIDs, sizeof(PIDs), &size))
	{
		return NULL;
	}

	size /= sizeof(DWORD);

	// Skip PID 0 (Idle Process) - always index 0?
	for (i = 1; i < size; i++)
	{
		HANDLE hProcess;
		char name[MAX_PATH];

		if (hProcess = OpenProcess(PROCESS_FLAGS, FALSE, PIDs[i]))
		{
			if (GetModuleBaseNameA(hProcess, NULL, name, MAX_PATH) == target_len && !lstrcmpiA(name, target))
			{
				*pid = PIDs[i];

				return hProcess;
			}

			CloseHandle(hProcess);
		}
	}

	return NULL;
}

void SetDEP(void)
{
	BOOL console = FALSE;
	DWORD flags = PROCESS_DEP_ENABLE;

	LPCSTR arg;

	int i;

	for (i = 1; i < __argc; i++)
	{
		arg = __argv[i];

		if (arg[0] == '-' || arg[0] == '/')
		{
			if ((arg[1] == 'd' || arg[1] == 'a') && arg[2] == '\0')
			{
				// Already used d or a
				if (flags != PROCESS_DEP_ENABLE)
				{
					Usage(console);

					return;
				}

				if (arg[1] == 'a')
				{
					flags |= PROCESS_DEP_DISABLE_ATL_THUNK_EMULATION;
				}
				else
				{
					flags = 0;
				}

				continue;
			}
			else if (arg[1] == 'c' && arg[2] == '\0')
			{
				console = TRUE;

				continue;
			}
		}

		break;
	}

	if (i == __argc - 1)
	{
		LPCSTR noAdmin = EnableDebugPrivilege() ? "" : "\n Note: Not running with full (Administrator) privileges";

		HANDLE hProcess;
		DWORD pid;

		char *end;

		arg = __argv[i];

		if (arg[0] <= '9' && arg[0] >= '0' && (pid = strtol(arg, &end, 10)) && *end == '\0')
		{
			hProcess = OpenProcess(PROCESS_FLAGS, FALSE, pid);
		}
		else
		{
			hProcess = OpenProcessByName(arg, &pid);
		}

		if (hProcess)
		{
			HANDLE hKernel32 = GetModuleHandleW(L"kernel32.dll");
			FARPROC SetProcessDEP = GetProcAddress(hKernel32, "SetProcessDEPPolicy");

			if (SetProcessDEP)
			{
				HANDLE procThread = CreateRemoteThread(hProcess, NULL, 0, (LPTHREAD_START_ROUTINE) SetProcessDEP, (LPVOID) flags, 0, NULL);
				LPCSTR status;

				char name[MAX_PATH];

				GetModuleBaseNameA(hProcess, NULL, name, MAX_PATH);

				if (procThread)
				{
					noAdmin = "";
					status = "called";

					CloseHandle(procThread);
				}
				else
				{
					status = "failed";
				}

				printf("SetProcessDEPPolicy %s in process %d (%s)%s", status, pid, name, noAdmin);
			}
			else
			{
				Message("SetProcessDEPPolicy function not found!", console);
			}

			CloseHandle(hProcess);
		}
		else
		{
			printf("Failed to find/open process%s", noAdmin);
		}
	}
	else
	{
		Usage(console);
	}
}

int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow)
{
	SetDEP();

	return 0;
}