TnbApiHooker.h

[詳解]

#pragma once
#include "TnbSimpleVector.h"
#include "TnbWindowsVersion.h"
 
 
 
 
 
//TNB Library
namespace TNB
{
 
 
 
template<class PFNC>
class CApiHookerT : CCopyImpossible
{
public:
    
    static PFNC Replace(LPCSTR lpszDllName, LPCSTR lpszFuncName, PFNC pFnc, HMODULE hModule = NULL)
    {
        CApiHookerT a;
        bool r = a.Hook(lpszDllName, lpszFuncName, pFnc, hModule);
        return r ? a() : NULL;
    }
 
 
    //----
 
 
    CApiHookerT(void) : m_pOrignalFnc(NULL), m_hModule(NULL)
    {
    }
 
    bool Hook(LPCSTR lpszDllName, LPCSTR lpszFuncName, PFNC pFnc, HMODULE hModule = NULL)
    {
        Unhook();
        CWindowsVersion wv;
        if ( ! wv.IsWow64() )
        {
            if ( hModule == NULL )
            {
                hModule = TNB::GetInstanceHandleByTnb(EI_Process);
            }
            m_modules.RemoveAll();
            m_pOrignalFnc = m_Replace(hModule, lpszDllName, lpszFuncName, pFnc);
            if ( m_pOrignalFnc != NULL )
            {
                m_hModule       = hModule;
                m_strDllName    = lpszDllName;
                m_strFuncName   = lpszFuncName;
                return true;
            }
        }
        return false;
    }
 
    bool Unhook(void)
    {
        if ( m_pOrignalFnc != NULL )
        {
            LPCSTR lpszFuncName = m_strFuncName;
            m_modules.RemoveAll();
            if ( m_Replace(m_hModule, m_strDllName, lpszFuncName, m_pOrignalFnc) == NULL )
            {
                return false;
            }
            m_pOrignalFnc = NULL;
            m_strDllName.Empty();
            m_strFuncName.Empty();
            m_hModule = NULL;
        }
        return true;
    }
 
    PFNC operator()(void)
    {
        return m_pOrignalFnc;
    }
 
private:
    PFNC                    m_pOrignalFnc;  
    HMODULE                 m_hModule;      
    CSimpleAscii            m_strDllName;   
    CSimpleAscii            m_strFuncName;  
    CSimpleVectorT<HMODULE> m_modules;      
 
    LPVOID m_DosHead(IMAGE_DOS_HEADER* pDos, DWORD offset)
    {
        return reinterpret_cast<LPVOID>(reinterpret_cast<DWORD>(pDos) + offset);
    }
 
    PFNC m_Replace(HMODULE hModule, LPCSTR lpszDllName, LPCSTR lpszFuncName, PFNC pFnc)
    {
        IMAGE_DOS_HEADER* pDOS = reinterpret_cast<IMAGE_DOS_HEADER*>(hModule);
        ASSERT( ! ::IsBadReadPtr(pDOS, sizeof(IMAGE_DOS_HEADER)) );
        if ( pDOS->e_magic != IMAGE_DOS_SIGNATURE )
        {
            return NULL;
        }
        IMAGE_NT_HEADERS* pPE = static_cast<IMAGE_NT_HEADERS*>(m_DosHead(pDOS, pDOS->e_lfanew));
        if ( pPE->Signature != IMAGE_NT_SIGNATURE )
        {
            return NULL;
        }
        IMAGE_DATA_DIRECTORY* pDataDir 
                    = reinterpret_cast<IMAGE_DATA_DIRECTORY*>(pPE->OptionalHeader.DataDirectory + IMAGE_DIRECTORY_ENTRY_IMPORT);
        if ( pDataDir->VirtualAddress == 0 || pDataDir->Size == 0 )
        {
            return NULL;
        }
        m_modules.Add(hModule);
        PFNC pResultFnc = NULL;
        //-- 下のDLLを検索
        // import DLL table
        IMAGE_IMPORT_DESCRIPTOR* pImpDsc = static_cast<IMAGE_IMPORT_DESCRIPTOR*>(m_DosHead(pDOS, pDataDir->VirtualAddress));
        for ( ; pImpDsc->FirstThunk != NULL; pImpDsc++ )
        {
            LPCSTR lpszName = static_cast<LPCSTR>(m_DosHead(pDOS, pImpDsc->Name));
            TTRACE1A("Found DLL = [%s]\n", lpszName);
            HMODULE h = ::GetModuleHandleA(lpszName);
            if ( h != NULL )
            {
                bool isFound = false;
                loop ( i, m_modules.GetSize() )
                {
                    if ( m_modules[i] == h )
                    {
                        isFound = true;
                        break;
                    }
                }
                if ( ! isFound )
                {
                    PFNC PF = m_Replace(h, lpszDllName, lpszFuncName, pFnc);
                    if ( PF != NULL )
                    {
                        if ( pResultFnc != NULL && pResultFnc != PF )
                        {
                            ASSERTLIB( false );
                        }
                        pResultFnc = PF;
                    }
                }
            }
        }
        //
        pImpDsc = static_cast<IMAGE_IMPORT_DESCRIPTOR*>(m_DosHead(pDOS, pDataDir->VirtualAddress));
        for ( ; pImpDsc->FirstThunk != NULL; pImpDsc++ )
        {
            LPCSTR lpszName = static_cast<LPCSTR>(m_DosHead(pDOS, pImpDsc->Name));
            TTRACE1A("Found DLL = [%s]\n", lpszName);
            if ( lpszDllName != NULL && stricmp(lpszName, lpszDllName) != 0 )
            {
                continue;
            }
            //-- 対象のDLL発見
            // address table
            IMAGE_THUNK_DATA* pIAT = static_cast<IMAGE_THUNK_DATA*>(m_DosHead(pDOS, pImpDsc->FirstThunk));
            // name table
            IMAGE_THUNK_DATA* pINT = static_cast<IMAGE_THUNK_DATA*>(m_DosHead(pDOS, pImpDsc->OriginalFirstThunk));
            for ( ; pIAT->u1.Function != NULL; pIAT++, pINT++ ) 
            {
                if ( IMAGE_SNAP_BY_ORDINAL(pINT->u1.Ordinal) )
                {
                    continue;
                }
                IMAGE_IMPORT_BY_NAME* pImpName = static_cast<IMAGE_IMPORT_BY_NAME*>(m_DosHead(pDOS, pINT->u1.AddressOfData));
                LPCSTR lpszName = reinterpret_cast<LPCSTR>(pImpName->Name);
                TTRACE1A("  - Found Func = [%s]\n", lpszName);
                if( strcmp(lpszName, lpszFuncName) == 0 ) 
                {
                    //-- 対象の関数発見
                    DWORD oldProt;
                    ::VirtualProtect(&pIAT->u1.Function, sizeof(DWORD), PAGE_READWRITE, &oldProt);
                    PFNC PF = reinterpret_cast<PFNC>(pIAT->u1.Function);
                    ASSERTLIB( PF != NULL );
                    pIAT->u1.Function = reinterpret_cast<DWORD>(pFnc);
                    DWORD dummy;
                    ::VirtualProtect(&pIAT->u1.Function, sizeof(DWORD), oldProt, &dummy);
                    if ( pResultFnc != NULL && pResultFnc != PF )
                    {
                        ASSERTLIB( false );
                    }
                    pResultFnc = PF;
                    break;
                }
            }
            if ( lpszDllName != NULL )
            {
                break;
            }
        }
        return pResultFnc;
    }
};
 
 
 
};//TNB