#include <Windows.h>
#include <Tlhelp32.h>
#include <iostream>

// 用于保存线程 TID 的数组
ULONG ThreadIds[100] = { 0 };

// 提权
void EnableDebugPriv()
{
	HANDLE hToken;
	LUID sedebugnameValue;
	TOKEN_PRIVILEGES tkp;

	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &hToken))
	{
		std::cout << "提权失败。" << std::endl;
		return;
	}

	if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &sedebugnameValue))
	{
		CloseHandle(hToken);
		std::cout << "提权失败。" << std::endl;
		return;
	}

	tkp.PrivilegeCount = 1;
	tkp.Privileges[0].Luid = sedebugnameValue;
	tkp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

	if (!AdjustTokenPrivileges(hToken, FALSE, &tkp, sizeof tkp, NULL, NULL))
	{
		std::cout << "提权失败。" << std::endl;
		CloseHandle(hToken);
	}
	else std::cout << "提权成功！" << std::endl;
}

// 根据进程名称获取 PID
// ProcessName 进程名称
ULONG GetProcessIdByProcessName(char * ProcessName)
{
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (INVALID_HANDLE_VALUE == hSnapshot)
	{
		return 0;
	}
	PROCESSENTRY32 pi;
	pi.dwSize = sizeof(PROCESSENTRY32); //第一次使用必须初始化成员
	BOOL bRet = Process32First(hSnapshot, &pi);
	while (bRet)
	{
		char Name[MAX_PATH] = { 0 };
		sprintf(Name, "%ls", pi.szExeFile);
		if (strcmp(Name, ProcessName) == 0)
		{
			return pi.th32ProcessID; // 返回PID
		}
		
		bRet = Process32Next(hSnapshot, &pi);
	}

	return 0;
}

// 根据 PID 获取获取所有相应的线程 TID
// ProcessId    进程 PID
// pThreadId   数组用于保存线程TID
// ThreadIdLen 用于返回数组实际长度
BOOL GetAllThreadIdByProcessId(ULONG ProcessId, ULONG * pThreadId, ULONG * ThreadIdLen)
{
	HANDLE hThreadSnap = INVALID_HANDLE_VALUE;
	THREADENTRY32 te32 = { sizeof(THREADENTRY32) };
	ULONG Number = 0;

	// 把所有线程拍一个快照
	hThreadSnap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (hThreadSnap == INVALID_HANDLE_VALUE)
		return(FALSE);

	// 在使用 Thread32First 前初始化 THREADENTRY32 的结构大小.
	te32.dwSize = sizeof(THREADENTRY32);

	// 现在获取系统线程列表, 并显示与指定进程相关的每个线程的信息	
	do {
		// 比对是否为该进程线程
		if (te32.th32OwnerProcessID == ProcessId)
		{
			// 是的话保存到线程数组中
			pThreadId[Number] = te32.th32ThreadID;
			Number++;
		}
	} while (Thread32Next(hThreadSnap, &te32));

	if (!Number)
		return FALSE;
	// 修改线程数量
	*ThreadIdLen = Number;
	return TRUE;
}

// APC注入
// ProcessName 要注入的进程名称
// DllName     要注入的DLL名称
BOOL ApcInjectDll(char * ProcessName, char * DllName)
{
	ULONG ProcessId = 0;
	ULONG IsFlag = TRUE;
	HANDLE hProcess = NULL, hThread = NULL;
	ULONG DllNameLen = strlen(DllName) + 1;
	PVOID pBaseAddr = NULL;
	ULONG dwRet = 0;
	PVOID pLoadLibraryA = NULL;
	ULONG ThreadIdsLen = 0;

	do 
	{
		// 根据进程名称获取 PID
		ProcessId = GetProcessIdByProcessName(ProcessName);
		if (0 >= ProcessId)
		{
			printf("ProcessId Error: %d\n", ProcessId);
			IsFlag = FALSE;
			break;
		}
		// 根据 PID 获取获取所有相应的线程 TID
		IsFlag = GetAllThreadIdByProcessId(ProcessId, ThreadIds, &ThreadIdsLen);
		if (!IsFlag)
		{
			printf("GetAllThreadIdByProcessId Error: %x\n", GetLastError());
			IsFlag = FALSE;
			break;
		}
		// 打开注入线程
		hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, ProcessId);
		if (hProcess == NULL)
		{
			
			printf("OpenProcess Error: %x\n", GetLastError());
			IsFlag = FALSE;
			break;
		}
		// 在注入的进程中申请空间
		pBaseAddr = 
			VirtualAllocEx(hProcess, NULL, DllNameLen, MEM_COMMIT | MEM_RESERVE, 
				PAGE_EXECUTE_READWRITE);
		if (pBaseAddr == NULL)
		{
			printf("VirtualAllocEx Error: %x\n", GetLastError());
			IsFlag = FALSE;
			break;
		}
		// 向申请的空间中写入 Dll 路径数据
		WriteProcessMemory(hProcess, pBaseAddr, DllName, DllNameLen, &dwRet);
		if (DllNameLen != dwRet)
		{
			printf("WriteProcessMemory Error: %x\n", GetLastError());
			IsFlag = FALSE;
			break;
		}
		// 获取 LoadLibraryA 地址
		pLoadLibraryA =
			GetProcAddress(GetModuleHandleA((char *)"kernel32.dll"), "LoadLibraryA");
		if (pLoadLibraryA == NULL)
		{
			printf("GetProcAddress Error: %x\n", GetLastError());
			IsFlag = FALSE;
			break;
		}
		// 遍历线程, 插入APC
		for (ULONG i = 0; i < ThreadIdsLen; i++)
		{
			// 打开线程
			hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, ThreadIds[i]);
			if (hThread != NULL)
			{
				// 插入 APC
				if (!QueueUserAPC((PAPCFUNC)pLoadLibraryA, hThread, (ULONG_PTR)pBaseAddr))
				{
					printf("QueueUserAPC Error: %x\n", GetLastError());
				}
				// 关闭线程句柄
				CloseHandle(hThread);
				hThread = NULL;
				printf("插入 APC 成功!\n");
			}
		}

	} while (FALSE);

	CloseHandle(hProcess);
	return IsFlag;
}

int main()
{
	// 提权
	EnableDebugPriv();
	// 注入我们想要执行的DLL
	ApcInjectDll((char *)"TestExe.exe", (char *)"./InjectDll.dll");
	system("pause");
	return 0;
}