gwangyi's Factory

#define UNICODE
#define _UNICODE

#include <windows.h>
#include <stdio.h>
#include <tchar.h>

LRESULT WINAPI HookProc(REFCLSID rclsid, LPUNKNOWN pUnkOuter, DWORD dwClsContext, REFIID riid, LPVOID * ppv);
HRESULT WINAPI hookedQueryInterface(LPUNKNOWN pThis, REFIID riid, void ** ppvObject);

BYTE backup[6] = {0, };
BYTE stub[6] = {0, };
BYTE qiStub[6] = {0, };
HRESULT WINAPI (*orgCoCreateInstance)(REFCLSID, LPUNKNOWN, DWORD, REFIID, LPVOID *);
HRESULT WINAPI (*pHookProc)(REFCLSID, LPUNKNOWN, DWORD, REFIID, LPVOID *) = HookProc;
HRESULT WINAPI (*pHookedQueryInterface)(LPUNKNOWN, REFIID, LPVOID *) = hookedQueryInterface;

struct MAP
{
    LPVOID orgAddress;
    BYTE stub[6];
} map[100];

int Initialize()
{
    LPVOID pRealAddress;
    HMODULE hModule;

    if(!(hModule = GetModuleHandle(TEXT("OLE32.DLL"))))
    {
        MessageBox(NULL, TEXT("OLE32.DLL is not found."), NULL, MB_OK);
        return FALSE;
    }
    if(!(pRealAddress = GetProcAddress(hModule, "CoCreateInstance")))
    {
        MessageBox(NULL, TEXT("Cannot found CoCreateInstance"), NULL, MB_OK);
        return FALSE;
    }

    orgCoCreateInstance = pRealAddress;
    memcpy(backup, orgCoCreateInstance, 6);
    *(WORD *)(stub) = 0x25ff;
    *(DWORD *)(&stub[2]) = (DWORD)&pHookProc;
    *(WORD *)(qiStub) = 0x25ff;
    *(DWORD *)(&qiStub[2]) = (DWORD)&pHookedQueryInterface;

    ZeroMemory(map, sizeof(map));

    return TRUE;
}

int Setup(BYTE stub[])
{
    DWORD flOldProtect, flNewProtect, flDontCare;
    MEMORY_BASIC_INFORMATION mbi;
    VirtualQuery((PVOID)orgCoCreateInstance, &mbi, sizeof(mbi));
    flNewProtect = mbi.Protect;
    flNewProtect &= ~(PAGE_READONLY | PAGE_EXECUTE_READ);
    flNewProtect |= (PAGE_READWRITE);
    VirtualProtect((PVOID)orgCoCreateInstance, 6, flNewProtect, &flOldProtect);
    memcpy((PVOID)orgCoCreateInstance, stub, 6);
    VirtualProtect((PVOID)orgCoCreateInstance, sizeof(PVOID), flOldProtect, &flDontCare);
    return TRUE;
}

int SetupQI(LPVOID obj)
{
    int i;
    DWORD flOldProtect, flNewProtect, flDontCare;
    MEMORY_BASIC_INFORMATION mbi;
    LPVOID qiAddr = **(LPVOID **)obj;

    for(i = 0; i < 100; i++)
    {
        if(map[i].orgAddress == qiAddr) return TRUE;
    }

    for(i = 0; i < 100; i++)
    {
        if(map[i].orgAddress == NULL) break;
    }
    if(i == 100) return FALSE;

    map[i].orgAddress = qiAddr;
    memcpy(map[i].stub, qiAddr, 6);

    VirtualQuery((PVOID)qiAddr, &mbi, sizeof(mbi));
    flNewProtect = (PAGE_READWRITE);
    if(!VirtualProtect((PVOID)qiAddr, 6, flNewProtect, &flOldProtect))
    {
        TCHAR buf[1000];
        _stprintf(buf, TEXT("Unable to hook QI of %p: you must set breakpoint manually at %p.\n GetLastError(): %x\n"), obj, qiAddr, GetLastError());
        OutputDebugString(buf);
        return TRUE;
    }
    memcpy((PVOID)qiAddr, qiStub, 6);
    VirtualProtect((PVOID)qiAddr, sizeof(PVOID), flOldProtect, &flDontCare);

    return TRUE;
}

int RestoreQI(LPVOID obj)
{
    int i;
    DWORD flOldProtect, flNewProtect, flDontCare;
    MEMORY_BASIC_INFORMATION mbi;
    LPVOID qiAddr = **(LPVOID **)obj;

    for(i = 0; i < 100; i++)
    {
        if(map[i].orgAddress == qiAddr) break;
    }
    if(i == 100) return FALSE;

    map[i].orgAddress = NULL;

    VirtualQuery((PVOID)qiAddr, &mbi, sizeof(mbi));
    flNewProtect = (PAGE_READWRITE);
    VirtualProtect((PVOID)qiAddr, 6, flNewProtect, &flOldProtect);
    memcpy((PVOID)qiAddr, map[i].stub, 6);
    VirtualProtect((PVOID)qiAddr, sizeof(PVOID), flOldProtect, &flDontCare);

    return TRUE;
}

int WINAPI DllMain(HANDLE h, DWORD dwReason, LPVOID lpReserved)
{
    if(dwReason == DLL_PROCESS_ATTACH)
    {
        MessageBox(NULL, TEXT("injected!!"), NULL, MB_OK);
        //TODO: initialize
        if(!Initialize()) return FALSE;

        if(Setup(stub))
            return TRUE;
        return FALSE;
    }
    else if(dwReason == DLL_PROCESS_DETACH)
    {
    }
    return TRUE;
}

BOOL isFilter(REFIID riid)
{
    TCHAR iid[40];
    StringFromGUID2(riid, iid, 40);
    return _tcsicmp(iid, TEXT("{56A86895-0AD4-11CE-B03A-0020AF0BA770}")) == 0;
}

HRESULT WINAPI HookProc(REFCLSID rclsid, LPUNKNOWN pUnkOuter, DWORD dwClsContext, REFIID riid, LPVOID * ppv)
{
    TCHAR out[1024];
    TCHAR clsid[40], iid[40];
    int ret;

    Setup(backup);
    ret = orgCoCreateInstance(rclsid, pUnkOuter, dwClsContext, riid, ppv);
    Setup(stub);

    if(isFilter(riid))
    {
        SetupQI(*ppv);
    }

    StringFromGUID2(rclsid, clsid, 40);
    StringFromGUID2(riid, iid, 40);
    _stprintf(out, TEXT("CoCreateInstance: CLSID: %s, IID: %s, object: %p\n"), clsid, iid, *ppv);
    OutputDebugString(out);
    return ret;
}

HRESULT WINAPI hookedQueryInterface(LPUNKNOWN pThis, REFIID riid, void ** ppvObject)
{
    HRESULT ret;
    TCHAR out[1024];
    TCHAR iid[40];
    RestoreQI(pThis);
    ret = pThis->lpVtbl->QueryInterface(pThis, riid, ppvObject);
    SetupQI(pThis);
    StringFromGUID2(riid, iid, 40);
    _stprintf(out, TEXT("QueryInterface: this: %p, IID: %s, object: %p, result: %08lX\n"), pThis, iid, *ppvObject, ret);
    OutputDebugString(out);
    if(*ppvObject != NULL)
        SetupQI(*ppvObject);

    return ret;
}