//修复IAT表
NTSTATUS RepairImportTable(PVOID pBaseAddress) //接受一个文件地址
{
	//DbgPrint(" go go go ");

	NTSTATUS status = STATUS_SUCCESS;
	PVOID pBaseAddressFun = NULL;
	PIMAGE_IMPORT_DESCRIPTOR pImportTableInfo = NULL;

	PIMAGE_NT_HEADERS pImageNT = (PIMAGE_NT_HEADERS)(((DWORD)pBaseAddress + ((PIMAGE_DOS_HEADER)pBaseAddress)->e_lfanew));

	//获取重定位表的RVA，在PIMAGE_DATA_DIRECTORY的第2个，每个PIMAGE_DATA_DIRECTORY位8个字节
	PIMAGE_DATA_DIRECTORY pRelocationTable = (PIMAGE_DATA_DIRECTORY)(pImageNT->OptionalHeader.DataDirectory);
	pRelocationTable++;

	if (pRelocationTable == NULL)
	{
		KdPrint(("获取DataDirectory数据失败！\n"));
		status = STATUS_UNSUCCESSFUL;
		return status;
	}

	//获取第一个导入表结构导入表结构
	pImportTableInfo = (PIMAGE_IMPORT_DESCRIPTOR)((DWORD)pBaseAddress + pRelocationTable->VirtualAddress);

	//循环导入表
	while (TRUE)
	{
		if (!pImportTableInfo->FirstThunk)
		{
			break;
		}

		//获取DLL名
		PCHAR dllName = (PCHAR)((DWORD)pBaseAddress + pImportTableInfo->Name);
		DWORD dwTimeDataStamp = pImportTableInfo->TimeDateStamp;
		PDWORD dwIATAddress = (PDWORD)(pImportTableInfo->FirstThunk + (DWORD)pBaseAddress);
		

		//获取函数名
		PIMAGE_THUNK_DATA pImageThunkData = (PIMAGE_THUNK_DATA)((DWORD)pBaseAddress + pImportTableInfo->OriginalFirstThunk);
		DWORD szNameAndSerialNumber = pImportTableInfo->OriginalFirstThunk;

		//获取模块基址
		pBaseAddressFun = GetDLLBaseAddress(dllName, dwTimeDataStamp);

		while (*((PDWORD)pImageThunkData))
		{
			//先判断是序号还是名称导入
			if ((DWORD)pImageThunkData->u1.AddressOfData & 0x80000000) //判断最高位是否为一
			{
				
			}
			else
			{
				//以名字导出
				PIMAGE_IMPORT_BY_NAME pImageImporotByName = (PIMAGE_IMPORT_BY_NAME)((PDWORD)((DWORD)pBaseAddress + pImageThunkData->u1.AddressOfData));

				//根据名字获取
				DWORD dwFunctionAddr = 0;
				dwFunctionAddr = (DWORD)GetMmGetSystemRoutineAddressEx((DWORD)pBaseAddressFun, (PSTR)((PUCHAR)pImageImporotByName->Name)); //根据函数名获取函数地址

				*dwIATAddress = dwFunctionAddr;   //修复IAT
			
				DbgPrint("名称导入(Hit/Name)：%d - %s , IAT : %x , 要修复的函数地址：%x\n", pImageImporotByName->Hint, (PUCHAR)pImageImporotByName->Name, (DWORD)(*dwIATAddress), (DWORD)dwFunctionAddr);
			}

			szNameAndSerialNumber += 4;
			pImageThunkData++;
			dwIATAddress++;
		}

		pImportTableInfo++;
		ZwUnmapViewOfSection(NtCurrentProcess(), pBaseAddressFun);
	}

	return status;
}


//获取dll的基址
PVOID GetDLLBaseAddress(PCHAR dllName, DWORD dwTimeDataStamp)
{
	PCHAR dllNampPath = ExAllocatePool(NonPagedPool, 1024);
	RtlZeroMemory(dllNampPath, 1024);
	RtlStringCbCopyA(dllNampPath, 1024, "\\??\\C:\\Windows\\System32\\");
	RtlStringCbCatA(dllNampPath, 1024, dllName);//追加字符串

	UNICODE_STRING strDllName;
	RtlInitUnicodeString(&strDllName, L"");  //初始化
	CHAR_TO_UNICODE_STRING(dllNampPath, &strDllName); //PCHAR to UNICODE


	PVOID pBaseAddressFun = NULL; //映射当前DLL模块，获取基址
	GetPEFileBaseAddressSection(strDllName, &pBaseAddressFun);  //映射当前DLL模块，获取基址

	DbgPrint("-------%s--------TimeDataStamp : %d------------%s----%wZ------%d--  \n", dllNampPath, dwTimeDataStamp, dllName, &strDllName, pBaseAddressFun);

	return pBaseAddressFun;
}

//获取导入dll模块的导出表，根据导出表获取导出的函数地址
PVOID GetMmGetSystemRoutineAddressEx(DWORD uModBase, CHAR* cSearchFnName) 
{
	IMAGE_DOS_HEADER* doshdr;
#ifdef AMD64
	IMAGE_OPTIONAL_HEADER64* opthdr;
#else
	IMAGE_OPTIONAL_HEADER32* opthdr;
#endif
	IMAGE_EXPORT_DIRECTORY* pExportTable;
	ULONG* dwAddrFns, * dwAddrNames;
	USHORT* dwAddrNameOrdinals;
	ULONG dwFnOrdinal, i;
	SIZE_T uFnAddr = 0;
	char* cFunName;
	doshdr = (IMAGE_DOS_HEADER*)uModBase;
	if (NULL == doshdr)
	{
		goto __exit;
	}
#ifdef AMD64
	opthdr = (IMAGE_OPTIONAL_HEADER64*)(uModBase + doshdr->e_lfanew + sizeof(ULONG) + sizeof(IMAGE_FILE_HEADER));
#else
	opthdr = (IMAGE_OPTIONAL_HEADER32*)(uModBase + doshdr->e_lfanew + sizeof(ULONG) + sizeof(IMAGE_FILE_HEADER));
#endif
	if (NULL == opthdr)
	{
		goto __exit;
	}
	pExportTable = (IMAGE_EXPORT_DIRECTORY*)(uModBase + opthdr->DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress);
	if (NULL == pExportTable)
	{
		goto __exit;
	}

	dwAddrFns = (ULONG*)(uModBase + pExportTable->AddressOfFunctions);
	dwAddrNames = (ULONG*)(uModBase + pExportTable->AddressOfNames);
	dwAddrNameOrdinals = (USHORT*)(uModBase + pExportTable->AddressOfNameOrdinals);
	for (i = 0; i < pExportTable->NumberOfNames; ++i)
	{
		cFunName = (char*)(uModBase + dwAddrNames[i]);
		if (!_strnicmp(cSearchFnName, cFunName, strlen(cSearchFnName)))
		{
			dwFnOrdinal = pExportTable->Base + dwAddrNameOrdinals[i] - 1;
			uFnAddr = uModBase + dwAddrFns[dwFnOrdinal];
			break;
		}
	}
__exit:
	return (PVOID)uFnAddr;
}