1 #ifndef _APITEST_IATHOOK_H
2 #define _APITEST_IATHOOK_H
3 
FindImportDescriptor(PBYTE DllBase,PCSTR DllName)4 static PIMAGE_IMPORT_DESCRIPTOR FindImportDescriptor(PBYTE DllBase, PCSTR DllName)
5 {
6     ULONG Size;
7     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor = RtlImageDirectoryEntryToData((HMODULE)DllBase, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &Size);
8     while (ImportDescriptor->Name && ImportDescriptor->OriginalFirstThunk)
9     {
10         PCHAR Name = (PCHAR)(DllBase + ImportDescriptor->Name);
11         if (!lstrcmpiA(Name, DllName))
12         {
13             return ImportDescriptor;
14         }
15         ImportDescriptor++;
16     }
17     return NULL;
18 }
19 
RedirectIat(HMODULE TargetDll,PCSTR DllName,PCSTR FunctionName,ULONG_PTR NewFunction,ULONG_PTR * OriginalFunction)20 static BOOL RedirectIat(HMODULE TargetDll, PCSTR DllName, PCSTR FunctionName, ULONG_PTR NewFunction, ULONG_PTR* OriginalFunction)
21 {
22     PBYTE DllBase = (PBYTE)TargetDll;
23     PIMAGE_IMPORT_DESCRIPTOR ImportDescriptor = FindImportDescriptor(DllBase, DllName);
24     if (ImportDescriptor)
25     {
26         // On loaded images, OriginalFirstThunk points to the name / ordinal of the function
27         PIMAGE_THUNK_DATA OriginalThunk = (PIMAGE_THUNK_DATA)(DllBase + ImportDescriptor->OriginalFirstThunk);
28         // FirstThunk points to the resolved address.
29         PIMAGE_THUNK_DATA FirstThunk = (PIMAGE_THUNK_DATA)(DllBase + ImportDescriptor->FirstThunk);
30         while (OriginalThunk->u1.AddressOfData && FirstThunk->u1.Function)
31         {
32             if (!IMAGE_SNAP_BY_ORDINAL32(OriginalThunk->u1.AddressOfData))
33             {
34                 PIMAGE_IMPORT_BY_NAME ImportName = (PIMAGE_IMPORT_BY_NAME)(DllBase + OriginalThunk->u1.AddressOfData);
35                 if (!lstrcmpiA((PCSTR)ImportName->Name, FunctionName))
36                 {
37                     DWORD dwOld;
38                     VirtualProtect(&FirstThunk->u1.Function, sizeof(ULONG_PTR), PAGE_EXECUTE_READWRITE, &dwOld);
39                     *OriginalFunction = FirstThunk->u1.Function;
40                     FirstThunk->u1.Function = NewFunction;
41                     VirtualProtect(&FirstThunk->u1.Function, sizeof(ULONG_PTR), dwOld, &dwOld);
42                     return TRUE;
43                 }
44             }
45             OriginalThunk++;
46             FirstThunk++;
47         }
48         skip("Unable to find the Import %s!%s\n", DllName, FunctionName);
49     }
50     else
51     {
52         skip("Unable to find the ImportDescriptor for %s\n", DllName);
53     }
54     return FALSE;
55 }
56 
RestoreIat(HMODULE TargetDll,PCSTR DllName,PCSTR FunctionName,ULONG_PTR OriginalFunction)57 static BOOL RestoreIat(HMODULE TargetDll, PCSTR DllName, PCSTR FunctionName, ULONG_PTR OriginalFunction)
58 {
59     ULONG_PTR old = 0;
60     return RedirectIat(TargetDll, DllName, FunctionName, OriginalFunction, &old);
61 }
62 
63  #endif // _APITEST_IATHOOK_H
64 
65