#include <debug.h>
#include "itsutils.h"
#include "ptrutils.h"
#include "kernelmisc.h"
#include "cever_deps.h"
#include "MemoryUsage.h"
ITSUTILS_API HRESULT STDAPICALLTYPE ITGetContext(
        DWORD cbInput, BYTE *pbInput,
        DWORD *pcbOutput, GetContextResult **ppbOutput,
        IRAPIStream *pStream)
{
    WCHAR *cmdline;

//    debug("ITGetContext: i=%08lx(%d)  o=%08lx(%08lx) s=%08lx\n", 
//            cbInput, pbInput, pcbOutput, ppbOutput, pStream);

    cmdline= GetCommandLine();

    GetContextResult *pOut=NULL;

    *pcbOutput= PTR_DIFF(pOut, pOut->wszCmdLine+wcslen(cmdline)+1);

    pOut= *ppbOutput= (GetContextResult*)LocalAlloc(LPTR, *pcbOutput);

    pOut->dwProcessId= GetCurrentProcessId();
    pOut->hProcess= GetCurrentProcess();
    pOut->dwThreadId= GetCurrentThreadId();
    pOut->hThread= GetCurrentThread();
    pOut->dwCallerTrust= CeGetCallerTrust();
    pOut->dwCurrentTrust= CeGetCurrentTrust();
    wcscpy(pOut->wszCmdLine, cmdline);

    KernelMode _km;
    CalcMemoryUsage(pOut->dwMemoryUsed, pOut->dwMemoryFree, pOut->dwKernelMemory);

    //calctotals();
    //calcfree();
    return 0;
}


ITSUTILS_API HRESULT STDAPICALLTYPE ITGetProcessList(
        DWORD cbInput, GetProcessListParams *pbInput,
        DWORD *pcbOutput, GetProcessListResult **ppbOutput,
        IRAPIStream *pStream)
{
    KernelMode _km;


    GetProcessListResult *pOut=NULL;
    *pcbOutput= 0;
    CeProcessList list;
    if (!GetProcessList(list))
        return ERROR_NOT_FOUND;
    *pcbOutput= PTR_DIFF(pOut,  pOut->pe+list.size());

    pOut= *ppbOutput= (GetProcessListResult*)LocalAlloc(LPTR, *pcbOutput);
    if (pOut==NULL) {
        error("ITGetProcessList: LocalAlloc(0x%x)", *pcbOutput);
        *pcbOutput= 0;
        return GetLastError();
    }

    pOut->nEntries= list.size();
    memcpy(pOut->pe, vectorptr(list), sizeof(CEPROCESSENTRY)*list.size());

    return 0;
}

//----------------------------------------------
class ModuleEntry {
public:
    ModuleEntry(const MODULEENTRY32& me) {
        wcsncpy(entry.szModuleName, me.szModule, MAX_PATH); 
        entry.dwMemoryBase= (DWORD)me.modBaseAddr;
    }
    ModuleEntry(const struct Module *mod) {
        ModuleInfo mi;
        FillModuleInfo(mi, mod);
        _snwprintf(entry.szModuleName, MAX_PATH, L"%hs", mi.name.c_str());
        entry.dwMemoryBase = mi.membase;
        entry.dwVBase = mi.csegbase;
        entry.dwDBase = mi.dsegbase;
        entry.dwUsage = mi.usagemask;
        entry.hLib    = (DWORD)mod;
    }

    CEMODULEENTRY entry;
};
typedef std::vector<ModuleEntry> ModuleEntryList;

// bool GetModuleList(ModuleEntryList& list)
// {
//     HANDLE hTH= CreateToolhelp32Snapshot(TH32CS_GETALLMODS|TH32CS_SNAPALL, 0);
//     if (hTH==INVALID_HANDLE_VALUE) {
//         error("CreateToolhelp32Snapshot");
//         return false;
//     }
// 
//     MODULEENTRY32 me;
//     me.dwSize= sizeof(MODULEENTRY32);
// 
//     int procnr=0;
//     if (Module32First(hTH, &me))
//     {
//         do {
//             list.push_back(ModuleEntry(me));
//         } while (Module32Next(hTH, &me));
//     }
//     if (GetLastError()!=ERROR_NO_MORE_FILES) {
//         error("Module32First/Next");
//         return false;
//     }
// 
//     CloseToolhelp32Snapshot(hTH);
// 
//     debug("found %d modules\n", list.size());
//     return true;
// }
bool GetMemoryModuleList(ModuleEntryList& list)
{
    struct Module *mod= (struct Module*)KData.aInfo[KINX_MODULES];

    while (mod) {
        list.push_back(ModuleEntry(mod));
        mod= mod->pMod;
    }
    return true;
}


ITSUTILS_API HRESULT STDAPICALLTYPE ITGetModuleList(
        DWORD cbInput, GetModuleListParams *pbInput,
        DWORD *pcbOutput, GetModuleListResult **ppbOutput,
        IRAPIStream *pStream)
{
    KernelMode _km;

    GetModuleListResult *pOut=NULL;

    ModuleEntryList modlist;

    bool bRes;
//    if (pbInput->bDirectRead)
        bRes= GetMemoryModuleList(modlist);
//  else
//      bRes= GetModuleList(modlist);


    if (!bRes)
        return GetLastError();

    *pcbOutput= PTR_DIFF(pOut,  pOut->me+modlist.size());
    pOut= *ppbOutput= (GetModuleListResult*)LocalAlloc(LPTR, *pcbOutput);
    if (pOut==NULL) {
        error("ITGetModuleList: LocalAlloc(0x%x)", *pcbOutput);
        *pcbOutput= 0;
        return GetLastError();
    }

    pOut->nEntries= modlist.size();
    memcpy(pOut->me, vectorptr(modlist), sizeof(CEMODULEENTRY)*modlist.size());

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

    return getProcessHandle(proc);
}
ITSUTILS_API HRESULT STDAPICALLTYPE ITGetProcessHandle(
        DWORD cbInput, GetProcessHandleParams *pbInput,
        DWORD *pcbOutput, GetProcessHandleResult **ppbOutput,
        IRAPIStream *pStream)
{
    *pcbOutput=sizeof(GetProcessHandleResult);
    *ppbOutput=(GetProcessHandleResult*)LocalAlloc(LPTR, *pcbOutput);

    KernelMode _km;

    (*ppbOutput)->hProc= GetProcessHandle(pbInput->wszProcessName);

    return 0;
}



ITSUTILS_API HRESULT STDAPICALLTYPE ITGetThreadTimes(
        DWORD cbInput, GetThreadTimesParams *pbInput,
        DWORD *pcbOutput, GetThreadTimesResult **ppbOutput,
        IRAPIStream *pStream)
{
    *pcbOutput= sizeof(GetThreadTimesResult);
    GetThreadTimesResult *pOut= *ppbOutput= (GetThreadTimesResult *)LocalAlloc(LPTR, *pcbOutput);

    KernelMode _km;

    DWORD res= 0;
    if (!GetThreadTimes(pbInput->hThread, &pOut->tCreate, &pOut->tExit, &pOut->tKernel, &pOut->tUser))
        res= GetLastError();

    return res;
}

ITSUTILS_API HRESULT STDAPICALLTYPE ITGetProcessUsageList(
        DWORD cbInput, GetProcessUsageListParams *pbInput,
        DWORD *pcbOutput, GetProcessUsageListResult **ppbOutput,
        IRAPIStream *pStream)
{
    KernelMode _km;

    GetProcessUsageListResult *pOut = *ppbOutput = NULL;
    *pcbOutput = 0;

    ProcSummaryInfoVector v;
    if (!GetProcessUsageList(v))
        return GetLastError() ? GetLastError() : ERROR_NOT_FOUND;
    *pcbOutput= PTR_DIFF(pOut, pOut->list+v.size());
    pOut = *ppbOutput= (GetProcessUsageListResult *)LocalAlloc(LPTR, *pcbOutput);

    memset(pOut, 0, sizeof(GetProcessUsageListResult));
    memcpy(pOut->list, &v[0], v.size()*sizeof(*pOut->list));

    pOut->nEntries= v.size();

    pOut->tQuery= GetTickCount();

    return 0;
}

ITSUTILS_API HRESULT STDAPICALLTYPE ITGetThreadUsageList(
        DWORD cbInput, GetThreadUsageListParams *pbInput,
        DWORD *pcbOutput, GetThreadUsageListResult **ppbOutput,
        IRAPIStream *pStream)
{
    KernelMode _km;

    GetThreadUsageListResult *pOut = *ppbOutput = NULL;
    *pcbOutput = 0;
    ThreadSummaryInfoVector v;
    if (!GetThreadUsageList(v, pbInput?pbInput->resolve_modulenames:true))
        return GetLastError() ? GetLastError() : ERROR_NOT_FOUND;
    *pcbOutput= PTR_DIFF(pOut, pOut->list+v.size());
    pOut = *ppbOutput= (GetThreadUsageListResult *)LocalAlloc(LPTR, *pcbOutput);

    memset(pOut, 0, sizeof(GetThreadUsageListResult));

    memcpy(pOut->list, &v[0], v.size()*sizeof(*pOut->list));

    pOut->nThreads= v.size();

    pOut->tQuery= GetTickCount();
    return 0;
}


// returns 0 when process not found, 1 when process not killed, 2 when process killed.
ITSUTILS_API HRESULT STDAPICALLTYPE ITTerminateProcess(
        DWORD cbInput, TerminateProcessParams *pbInput,
        DWORD *pcbOutput, BYTE **ppbOutput,
        IRAPIStream *pStream)
{
    *ppbOutput=NULL;
    *pcbOutput=0;

    KernelMode _km;

    HANDLE hProc;
    if (pbInput->dwProcessHandle)
        hProc= OpenProcess(0, 0, pbInput->dwProcessHandle);
    else
        hProc= GetProcessHandle(pbInput->wszProcessName);

    if (hProc==NULL || hProc==INVALID_HANDLE_VALUE) {
        return ERROR_NOT_FOUND;
    }

    DWORD res= 0;
    if (!TerminateProcess(hProc, 0))
        res= GetLastError();
    else if (pbInput->bWait)
        WaitForSingleObject(hProc, INFINITE);
    CloseHandle(hProc);

    return res;
}

// returns 0 when process not found, 1 when process not killed, 2 when process killed.
ITSUTILS_API HRESULT STDAPICALLTYPE ITTerminateThread(
        DWORD cbInput, TerminateThreadParams *pbInput,
        DWORD *pcbOutput, BYTE **ppbOutput,
        IRAPIStream *pStream)
{
    *ppbOutput=NULL;
    *pcbOutput=0;

    HANDLE hThread= (HANDLE)pbInput->dwThreadHandle;

    if (hThread==NULL || hThread==INVALID_HANDLE_VALUE) {
        return ERROR_NOT_FOUND;
    }

    DWORD res= 0;
    if (!TerminateThread(hThread, 0))
        res= GetLastError();
    CloseHandle(hThread);

    return res;
}

// returns 0 when process not found, 1 when process not killed, 2 when process killed.
ITSUTILS_API HRESULT STDAPICALLTYPE ITWaitForProcess(
        DWORD cbInput, WaitForProcessParams *pbInput,
        DWORD *pcbOutput, BYTE **ppbOutput,
        IRAPIStream *pStream)
{
    *ppbOutput=NULL;
    *pcbOutput=0;

    KernelMode _km;

    HANDLE hProc = OpenProcess(0, 0, pbInput->dwProcessHandle);

    if (hProc==NULL || hProc==INVALID_HANDLE_VALUE) {
        return ERROR_NOT_FOUND;
    }

    DWORD res= 0;
    if (WAIT_OBJECT_0!=WaitForSingleObject(hProc, INFINITE)) {
        res= GetLastError();
    }
    else {
        GetExitCodeProcess(hProc, &res);
    }
    CloseHandle(hProc);

    return res;
}