#include "ApiHook.h"
#include "WindowsVersionHelp.h"
#include <tlhelp32.h>
#include <stdio.h>
#include <delayimp.h>


namespace api_hook {
	void replaceAddress(IMAGE_THUNK_DATA* importAddressTable, const void* pfnCurrent, const void* pfnNew) {
        for (; importAddressTable->u1.Function; ++importAddressTable) {
            void** ppfn = reinterpret_cast<void**>(&importAddressTable->u1.Function);
            if (*ppfn != pfnCurrent) {
                continue;
            }
            DWORD dwDummy = 0;
            VirtualProtect(ppfn, sizeof(ppfn), PAGE_EXECUTE_READWRITE, &dwDummy);
            // DLL code will be CoW
            WriteProcessMemory(GetCurrentProcess(), ppfn, &pfnNew, sizeof(pfnNew), NULL);
            puts("  +++++ REPLACE OK +++++");
            return;
        }
	}


    void delay(
        const char* pszModuleName,
        const void* pfnCurrent,
        const void* pfnNew,
        HMODULE hmodCaller) {

        ULONG ulSize = 0;

        ImgDelayDescr* delayDescr = 
            reinterpret_cast<ImgDelayDescr*>(
                dll.ImageDirectoryEntryToData(
                    hmodCaller, TRUE, IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT, &ulSize));

        if (!delayDescr) {
            return;
        }

        printf("  delay import\n");
        char* const base = reinterpret_cast<char* const>(hmodCaller);
        for (; delayDescr->grAttrs; ++delayDescr) {
            char* const name = base + delayDescr->rvaDLLName;
            printf("    name=%08p:%s\n", base, name);
            if (lstrcmpiA(name, pszModuleName) == 0) {
                break;
            }
        }

        if (!delayDescr->grAttrs) {
            return;
        }

		IMAGE_THUNK_DATA* importAddressTable = 
			reinterpret_cast<IMAGE_THUNK_DATA*>(base + delayDescr->rvaIAT);

		replaceAddress(importAddressTable, pfnCurrent, pfnNew);
    }

    void normal(
        const char* pszModuleName,
        const void* pfnCurrent,
        const void* pfnNew,
        HMODULE hmodCaller) {

        printf("DLL: %s\n", pszModuleName);

        ULONG ulSize = 0;

        IMAGE_IMPORT_DESCRIPTOR* pImportDesc = 
            reinterpret_cast<IMAGE_IMPORT_DESCRIPTOR*>(
                dll.ImageDirectoryEntryToData(
                    hmodCaller, TRUE, IMAGE_DIRECTORY_ENTRY_IMPORT, &ulSize));

        if (!pImportDesc) {
            return;
        }

        for (; pImportDesc->Characteristics; ++pImportDesc) {
            char* base = reinterpret_cast<char*>(hmodCaller);
            ptrdiff_t off = pImportDesc->Name;
            char* pszModName = base + off;
            printf("  base=0x%p, off=0x%x pszModName=%s\n", base, off, pszModName);
            if (lstrcmpiA(pszModName, pszModuleName) == 0) {
                break;
            }
        }

        if (!pImportDesc->Characteristics) {
            return;
        }

        IMAGE_THUNK_DATA* importAddressTable = 
            reinterpret_cast<IMAGE_THUNK_DATA*>(
                reinterpret_cast<char*>(hmodCaller) + pImportDesc->FirstThunk);

		replaceAddress(importAddressTable, pfnCurrent, pfnNew);
    }

    /**
     * ひとつのモジュールに対してAPIフックを行う関数
     * @see http://ruffnex.oc.to/kenji/text/api_hook/
     */
    template <>
    void ReplaceIATEntryInOneMod<const void, const void>(
        const char* pszModuleName,
        const void* pfnCurrent,
        const void* pfnNew,
        HMODULE hmodCaller)
    {
        normal(pszModuleName, pfnCurrent, pfnNew, hmodCaller);
        delay(pszModuleName, pfnCurrent, pfnNew, hmodCaller);
    }

    /**
     * すべてのモジュールに対してAPIフックを行う関数
     * @see http://ruffnex.oc.to/kenji/text/api_hook/
     */
    template <>
    void ReplaceIATEntryInAllMods<const void, const void>(
        const char* pszModuleName,
        const void* pfnCurrent,
        const void* pfnNew)
    {
        // モジュールリストを取得
        HANDLE hModuleSnap = CreateToolhelp32Snapshot(
                                 TH32CS_SNAPMODULE, GetCurrentProcessId());
        if(hModuleSnap == INVALID_HANDLE_VALUE) {
            return;
        }
    
        MODULEENTRY32 me = {};
        me.dwSize = sizeof(me);
        BOOL bModuleResult = Module32First(hModuleSnap, &me);
        // それぞれのモジュールに対してReplaceIATEntryInOneModを実行
        while (bModuleResult) {
            wprintf(L"ReplaceIATEntryInAllMods: %s\n", me.szExePath);
            ReplaceIATEntryInOneMod(pszModuleName, pfnCurrent, pfnNew, me.hModule);
            bModuleResult = Module32Next(hModuleSnap, &me);
        }
        CloseHandle(hModuleSnap);
    }
}