#include <windows.h>
#include <string>
#include "stringutils.h"
#include "debug.h"
#include "cenk.h"
#include "cever_deps.h"

// pregutl -s hklm\init :Launch09='loadhook.exe' :Depend10=hex:09,00

// on sm2003  loading from init somehow does not work.
// 
// prun loadhook.exe -load apihookw32a.dll 20
// prun loadhook.exe -load apihookw32h.dll  7
// prun loadhook.exe -load apihookw32s.dll  0
// prun loadhook.exe -load apihookwmgr.dll 17
// prun loadhook.exe -load apihookgdi.dll 16


class BaseError {
protected:
    std::Wstring msg_;
    va_list ap_;
public:
    BaseError(WCHAR *fmt, va_list ap)
    {
        int desired_length= 1024;
        msg_.resize(desired_length);
        int printedlength= _vsnwprintf(&msg_[0], msg_.size(), fmt, ap);

        va_end(ap);

        if (printedlength!=-1)
            msg_.resize(printedlength);
    }
    const WCHAR* c_str() const
    {
        return msg_.c_str();
    }
};
class SystemError : public BaseError {
private:
    DWORD dwErrorCode;
    std::Wstring err_;
public:
    SystemError(WCHAR *fmt, ...)
        : BaseError(fmt, (va_start(ap_, fmt), ap_))
        , dwErrorCode(GetLastError())
    {
        WCHAR* msgbuf;
        int rc= FormatMessage(
                FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
                NULL, dwErrorCode, 0, (WCHAR*) &msgbuf, 0, NULL);
        if (rc)
        {
            err_= msg_ + L": " + msgbuf;
            LocalFree(msgbuf);
        }
        else {
            err_= msg_ + L": " + ToWString(stringformat("%08lx", dwErrorCode));
        }
    }
    const WCHAR* c_str() const
    {
        return err_.c_str();
    }
    ~SystemError() { printf("SYSERROR: %08lx  %s\n", dwErrorCode, ToString(err_).c_str()); }
};

class win32api {
public:
    static HMODULE GetModuleHandle(const std::Wstring& modulename)
    {
        HMODULE hMod= ::GetModuleHandle(modulename.c_str());
        if (hMod==NULL || hMod==INVALID_HANDLE_VALUE)
            throw SystemError(L"GetModuleHandle(%ls)", modulename.c_str());
        return hMod;
    }
    static HMODULE LoadLibrary(const std::Wstring& modulename)
    {
        HMODULE hMod= ::LoadLibrary(modulename.c_str());
        if (hMod==NULL || hMod==INVALID_HANDLE_VALUE)
            throw SystemError(L"LoadLibrary(%ls)", modulename.c_str());
        return hMod;
    }
    static HMODULE LoadKernelLibrary(const std::Wstring& modulename)
    {
        HMODULE hMod= (HMODULE)::LoadKernelLibrary(modulename.c_str());
        if (hMod==NULL || hMod==INVALID_HANDLE_VALUE)
            throw SystemError(L"LoadKernelLibrary(%ls)", modulename.c_str());
        return hMod;
    }
    static LPVOID GetProcAddress(HMODULE hMod, const std::Wstring& procname)
    {
        LPVOID pfn= ::GetProcAddress(hMod, procname.c_str());
        if (pfn==NULL)
            throw SystemError(L"GetProcAddress(%08lx, %ls)", hMod, procname.c_str());
        return pfn;
    }
};

HANDLE GetProcessHandle(const WCHAR *wszProcessName)
{
    PROCESS *proc= FindProcessForName(wszProcessName);
    if (proc==NULL)
        return INVALID_HANDLE_VALUE;

    return getProcessHandle(proc);
}
DWORD RemoteLoadLibrary(HANDLE hProcess, const std::Wstring& dllname)
{
    if (hProcess==NULL)
        return (DWORD)win32api::LoadKernelLibrary(dllname);

    CALLBACKINFO cbi;
    cbi.hProc= hProcess;

    HMODULE hCoredll= win32api::GetModuleHandle(L"coredll.dll");
    LPVOID pfnLoadLib= win32api::GetProcAddress(hCoredll, L"LoadLibraryW");
    cbi.pfn= (FARPROC)MapPtrToProcess(pfnLoadLib, cbi.hProc);
    cbi.pvArg0= MapPtrToProcess((void*)dllname.c_str(), GetCurrentProcess());

    debug("loadhook: proc=%08lx coredll=%08lx pfnloadlib=%08lx->%08lx  [%s]\n", hProcess, hCoredll, pfnLoadLib, cbi.pfn, ToString(dllname).c_str());

    return PerformCallBack4(&cbi, 0, 0, 0);
}
DWORD RemoteFreeLibrary(HANDLE hProcess, DWORD hModule)
{
    if (hProcess==NULL)
        return FALSE;
    CALLBACKINFO cbi;
    cbi.hProc= hProcess;

    HMODULE hCoredll= win32api::GetModuleHandle(L"coredll.dll");
    LPVOID pfnFreeLib= win32api::GetProcAddress(hCoredll, L"FreeLibrary");
    cbi.pfn= (FARPROC)MapPtrToProcess(pfnFreeLib, cbi.hProc);
    cbi.pvArg0= (LPVOID)hModule;

    debug("loadhook: proc=%08lx coredll=%08lx pfnloadlib=%08lx->%08lx\n", hProcess, hCoredll, pfnFreeLib, cbi.pfn);

    return PerformCallBack4(&cbi, 0, 0, 0);
}

// used to start/stop apihook
DWORD RemoteCallHookApi(const std::string& dllname, HANDLE hProcess, const std::Wstring& apiname, DWORD p0, DWORD p1, DWORD p2, DWORD p3)
{
    if (hProcess==NULL) {
        HMODULE hKmod= win32api::LoadKernelLibrary(ToWString(dllname));
        if (hKmod==NULL) {
            error("RemoteCallHookApi::LoadKernelLibrary");
            return 0;
        }
        typedef DWORD(*callback_t)(DWORD p0, DWORD p1, DWORD p2, DWORD p3);
        callback_t pfn= (callback_t)win32api::GetProcAddress(hKmod, apiname);
        return pfn(p0,p1,p2,p3);
    }
    CALLBACKINFO cbi;
    cbi.hProc= hProcess;

    static HMODULE hLocalHookMod;
    if (hLocalHookMod==NULL) {
        hLocalHookMod= win32api::LoadLibrary(ToWString(dllname));
        debug("loadhook: local=%08lx\n", hLocalHookMod);
    }
    LPVOID pfn= win32api::GetProcAddress(hLocalHookMod, apiname);

    cbi.pfn= (FARPROC)MapPtrToProcess(pfn, cbi.hProc);
    cbi.pvArg0= (void*)p0;

    debug("loadhook: remotecall proc=%08lx pfn=%08lx -> %08lx %s %ls\n", hProcess, pfn, cbi.pfn, dllname.c_str(), apiname.c_str());

    return PerformCallBack4(&cbi, p1, p2, p3);
}

void loadhook(const std::string& dllname, HANDLE hProc)
{
    try {
        DWORD hRemoteHookMod= RemoteLoadLibrary(hProc, ToTString(dllname).c_str());
        debug("loadhook: remote=%08lx\n", hRemoteHookMod);

        DWORD res= RemoteCallHookApi(dllname, hProc, L"starthook", 0,0,0,0);
        debug("loadhook: starthook: %08lx\n", res);

        Sleep(1000);
        DWORD count= RemoteCallHookApi(dllname, hProc, L"callcount", 0,0,0,0);
        debug("loadhook: count: %08lx\n", count);
    }
    catch (SystemError e) {
        debug("loadhook: ERROR: %ls\n", e.c_str());
    }
}
void unloadhook(const std::string& dllname, HANDLE hProc)
{
    try {
        DWORD count= RemoteCallHookApi(dllname, hProc, L"callcount", 0,0,0,0);
        debug("unloadhook: count: %08lx\n", count);

        DWORD res= RemoteCallHookApi(dllname, hProc, L"endhook", 0,0,0,0);
        debug("unloadhook: endhook: %08lx\n", res);

        HMODULE hApiHookModule= win32api::GetModuleHandle(ToTString(dllname).c_str());
        debug("unloadhook: hh= %08lx\n", hApiHookModule);
        res= RemoteFreeLibrary(hProc, (DWORD)hApiHookModule);
        debug("unloadhook: freeremote=%08lx\n", res);
    }
    catch (SystemError e) {
        debug("unloadhook: ERROR: %ls\n", e.c_str());
    }
}
// PUserKData      c:/local/wince420/PUBLIC/COMMON/SDK/INC/kfuncs.h
// KDataStruct     c:/local/wince420/PRIVATE/WINCEOS/COREOS/NK/INC/nkarm.h
// KINX_APISETS    c:/local/wince420/PUBLIC/COMMON/OAK/INC/pkfuncs.h
// SH_FILESYS_APIS c:/local/wince420/PUBLIC/COMMON/SDK/INC/kfuncs.h
// CINFO           c:/local/wince420/PRIVATE/WINCEOS/COREOS/NK/INC/kernel.h
// PROCESS         c:/local/wince420/PRIVATE/WINCEOS/COREOS/NK/INC/kernel.h
// HT_FILE         c:/local/wince420/PUBLIC/COMMON/OAK/INC/psyscall.h
int WINAPI WinMain( HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
    DebugSetLogfile("hook.log");

    StringList args;
    if (!SplitString(ToString(lpCmdLine), args, false))
    {
        error("Error in commandline");
        return 1;
    }
    bool bDoUnload= false;
    bool bDoLoad= false;
    int apinr=0;
    std::string dllname;
    std::string procname;
    for (StringList::iterator i= args.begin() ; i!=args.end() ; ++i)
    {
        std::string& arg= *i;
        if (arg[0]=='-') switch(arg[1])
        {
            case 'u': bDoUnload= true; break;
            case 'l': bDoLoad= true; break;
            case 'n': procname= arg; break;
        }
        else {
            if (dllname.empty())
                dllname= arg;
            else
                apinr= strtoul(arg.c_str(), 0, 0);
        }
    }

    // setkmode+setprocperms is only nescesary under wm2005
    // for ppc2003 loadhook works fine without.
    BOOL bMode = SetKMode(TRUE);
    DWORD dwPerm = SetProcPermissions(0xFFFFFFFF);

    debug("loadhook: %08lx:%08lx started hProc=%08lx hThread=%08lx: %ls\n", hCurProc, hCurThread, hCurProc, hCurThread, lpCmdLine);

    CINFO **apisets= (CINFO **)KData.aInfo[KINX_APISETS];
    debug("loadhook: apisets=%08lx\n", apisets);

    HANDLE hProc;
    if (apinr) {
        CINFO *api= apisets[apinr];
        hProc= api->pServer?api->pServer->hProc:0;
        debug("loadhook: api[%d]: %08lx %c%c%c%c %d methods, svr: %08lx ( procobj=%08lx hproc=%08lx )\n", apinr, api,
                api->acName[0], api->acName[1], api->acName[2], api->acName[3]?api->acName[3]:' ',
                api->cMethods, api->pServer, api->pServer, hProc);
    }
    else if (!procname.empty()) {
        hProc= GetProcessHandle(ToWString(procname).c_str());
    }

    if (bDoUnload)
        unloadhook(dllname, hProc);
    else if (bDoLoad) {
        loadhook(dllname, hProc);
        // calling SignalStarted to be able to run loadhook from HKLM\Init
        DWORD dwInitValue= _wtol(lpCmdLine);
        if (dwInitValue)
            SignalStarted(dwInitValue);
    }

    SetProcPermissions(dwPerm);
    SetKMode(bMode);

    debug("loadhook: ended\n");
    return 0;
}