#include <windows.h>

// 定義API掛接項結構
typedef struct _HOOK_ITEM {
 DWORD dwAddr ;   // IAT項所在地址
 DWORD dwOldValue ;  // IAT項的原始函式地址
 DWORD dwNewValue ;  // IAT項的新函式地址
} HOOK_ITEM, *PHOOK_ITEM ;
HOOK_ITEM HookItem = {0} ; // 定義IAT項，用於儲存MessageBoxA的IAT項資訊

// 定義MessageBoxA函式原型
typedef int (WINAPI* PFNMessageBoxA)( HWND hWnd, LPCSTR lpText, LPCSTR lpCaption, UINT uType ) ;

// 定義重定向API的實現函式
BOOL WINAPI RedirectApi ( PCHAR pDllName, PCHAR pFunName, DWORD dwNewProc, PHOOK_ITEM pItem ) ;

// 自定義的MessageBoxA函式
// 實現對原始MessageBoxA的輸入、輸出引數的監控，甚至是取消呼叫
int WINAPI NEW_MessageBoxA( HWND hWnd, LPCSTR lpText, LPCSTR lpCaption, UINT uType )
{
 // 此處可以觀察/修改呼叫引數，甚至可以取消呼叫直接返回。
 // ……

 // 取得原函式地址
 PFNMessageBoxA pfnMessageBoxA = (PFNMessageBoxA)HookItem.dwOldValue ;

 // 輸出測試資訊，
 // 如果這裡直接呼叫MessageBoxA，就進入無限迴圈
 pfnMessageBoxA ( hWnd, "這是API重定向過程的訊息框", "測試", 0 ) ;

 // 呼叫原函式
 int ret = pfnMessageBoxA ( hWnd, lpText, lpCaption, uType ) ;

 // 此處可以檢視/修改呼叫原函式的返回值
 // ……

 return ret ;
}

int WinMain( HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, int nCmdShow )
{
 // 重定向API
 if ( !RedirectApi ( "USER32.dll", "MessageBoxA", (DWORD)NEW_MessageBoxA, &HookItem ) )
  OutputDebugStringA ( "RedirectApi failed!" ) ;
 else
  OutputDebugStringA ( "RedirectApi success!" ) ;
 
 MessageBoxA ( 0, "正常訊息框", "測試", 0 ) ;
 return 0 ;
}

// 實現重定向API
// 引數pDllName:目標API所在的DLL名稱
// 引數pFunName:目標API名稱
// 引數dwNewProc:自定義的函式地址
// 引數pItem:用於儲存IAT項資訊
BOOL WINAPI RedirectApi ( PCHAR pDllName, PCHAR pFunName, DWORD dwNewProc, PHOOK_ITEM pItem )
{
 // 檢查引數是否合法
 if ( pDllName == NULL || pFunName == NULL || !dwNewProc || !pItem )
  return FALSE ;

 // 檢測目標模組是否存在
 char szTempDllName[256] = {0} ;
 DWORD dwBaseImage = (DWORD)GetModuleHandle(NULL) ;
 if ( dwBaseImage == 0 )
  return FALSE ;

 // 取得PE檔案頭資訊指標
 PIMAGE_DOS_HEADER   pDosHeader = (PIMAGE_DOS_HEADER)dwBaseImage ;
 PIMAGE_NT_HEADERS   pNtHeader = (PIMAGE_NT_HEADERS)(dwBaseImage + (pDosHeader->e_lfanew)) ;
 PIMAGE_OPTIONAL_HEADER32 pOptionalHeader = &(pNtHeader->OptionalHeader) ;
 PIMAGE_SECTION_HEADER  pSectionHeader = (PIMAGE_SECTION_HEADER)((DWORD)pNtHeader + 0x18 + pNtHeader->FileHeader.SizeOfOptionalHeader ) ;

 // 遍歷匯入表
 PIMAGE_THUNK_DATA pThunk, pIAT ;
 PIMAGE_IMPORT_DESCRIPTOR pIID = (PIMAGE_IMPORT_DESCRIPTOR)(dwBaseImage+pOptionalHeader->DataDirectory[1].VirtualAddress ) ;
 while ( pIID->FirstThunk )
 {
  // 檢測是否目標模組
  if ( strcmp ( (PCHAR)(dwBaseImage+pIID->Name), pDllName ) )
  {
   pIID++ ;
   continue ;
  }

  pIAT = (PIMAGE_THUNK_DATA)( dwBaseImage + pIID->FirstThunk ) ;
  if ( pIID->OriginalFirstThunk )
   pThunk = (PIMAGE_THUNK_DATA)( dwBaseImage + pIID->OriginalFirstThunk ) ;
  else
   pThunk = pIAT ;

  // 遍歷IAT
  DWORD dwThunkValue = 0 ;
  while ( ( dwThunkValue = *((DWORD*)pThunk) ) != 0 )
  {
   if ( ( dwThunkValue & IMAGE_ORDINAL_FLAG32 ) == 0 )
   {
    // 檢測是否目標函式
    if ( strcmp ( (PCHAR)(dwBaseImage+dwThunkValue+2), pFunName ) == 0 )
    {
     // 填充函式重定向資訊
     pItem->dwAddr  = (DWORD)pIAT ;
     pItem->dwOldValue = *((DWORD*)pIAT) ;
     pItem->dwNewValue = dwNewProc;

     // 修改IAT項
     DWORD dwOldProtect = 0 ;
     VirtualProtect ( pIAT, 4, PAGE_READWRITE, &dwOldProtect ) ;
     *((DWORD*)pIAT) = dwNewProc ;
     VirtualProtect ( pIAT, 4, PAGE_READWRITE, &dwOldProtect ) ;
     return TRUE ;
    }
   }

   pThunk ++ ;
   pIAT ++ ;
  }

  pIID ++ ;
 }

 return FALSE ;
}