#include <windows.h>
#include <debug.h>
#include "itsutils.h"
#include "kernelmisc.h"
#include "ptrutils.h"
#define MEM_PAGE_SIZE 4096

#ifdef WITHDEBUG
#define memdebug debug
#else
#define memdebug while(0)
#endif
bool ReadMemory(DWORD dwStart, BYTE *buf, DWORD dwLength, DWORD *pdwCopied, DWORD nDataAccess)
{
    DWORD nCopied= 0;
    DWORD addr= dwStart;
    bool bIgnoreErrors= (nDataAccess&8)!=0;
    memdebug("readmem(a=%08lx, o=%08lx l=%08lx)\n", nDataAccess, dwStart, dwLength);

    __try {
    if ((nDataAccess&7)==0) {
        while (nCopied < dwLength)
        {
            DWORD nChunkSize= std::min(MEM_PAGE_SIZE-(addr&(MEM_PAGE_SIZE-1)), dwLength-nCopied);
            //debug("chunk(%08lx, %08lx) = %08lx\n", addr, nChunkSize, dwLength-nCopied);

            if (IsBadReadPtr((void*)addr, nChunkSize)) {
                if (bIgnoreErrors)
                    memset(&buf[nCopied], 0, nChunkSize);
                else
                    return false;
            }
            else
                memcpy(&buf[nCopied], (void*)addr, nChunkSize);
            nCopied += nChunkSize;
            addr += nChunkSize;
        }
        //debug("readmem(%08lx,%08lx) done: nc=%08lx ad=%08lx\n", dwStart, dwLength, nCopied, addr);
    }
    else if (IsBadReadPtr((void*)addr, dwLength)) {
        // no direct access.
        debug("readmem: no direct access\n");
    }
    else if ((nDataAccess&7)==1) {
        while (nCopied < dwLength)
        {
            *(BYTE*)(buf+nCopied)= *(BYTE*)addr;
            nCopied++;
            addr++;
        }
    }
    else if (addr&1) {
        // non-word aligned address not allowed
        debug("readmem: non-dword aligned address not allowed\n");
    }
    else if ((nDataAccess&7)==2) {
        while (nCopied+1 < dwLength)
        {
            *(WORD*)(buf+nCopied)= *(WORD*)addr;
            nCopied+=2;
            addr+=2;
        }
    }
    else if (addr&3) {
        // non-dword aligned address not allowed
        debug("readmem: non-dword aligned address not allowed\n");
    }
    else if ((nDataAccess&7)==4) {
        while (nCopied+3 < dwLength)
        {
            *(DWORD*)(buf+nCopied)= *(DWORD*)addr;
            nCopied+=4;
            addr+=4;
        }
    }

    *pdwCopied= nCopied;

    return true;
    } __except(EXCEPTION_EXECUTE_HANDLER) {
        memdebug("EXCEPTION, access=%x, addr=%08x, copied=%08x, dstbuf=%08x\n", nDataAccess, addr, nCopied, buf);
        return false;
    }

}
bool SafeReadProcMem(HANDLE hProc, DWORD dwStart, BYTE *buf, DWORD dwLength, DWORD *pdwCopied, bool bIgnoreErrors)
{
    DWORD nCopied= 0;
    DWORD addr= dwStart;

    memdebug("saferead(a=%08lx l=%08lx)\n", dwStart, dwLength);
    while (nCopied < dwLength)
    {
        DWORD nChunkSize= std::min(MEM_PAGE_SIZE-(addr&(MEM_PAGE_SIZE-1)), dwLength-nCopied);

        DWORD nRead;
        if (!ReadProcessMemory(hProc, (void*)addr, &buf[nCopied], nChunkSize, &nRead)) {
            if (bIgnoreErrors)
                memset(&buf[nCopied], 0, nChunkSize);
            else
                return false;
        }
        nCopied += nChunkSize;
        addr += nChunkSize;
    }
    *pdwCopied= nCopied;

    return true;
}

bool WriteMemory(DWORD dwStart, const BYTE *buf, DWORD dwLength, DWORD *pdwCopied, DWORD nDataAccess)
{
    DWORD nCopied= 0;
    DWORD addr= dwStart;
    bool bIgnoreErrors= (nDataAccess&8)!=0;

    if ((nDataAccess&7)==0) {
        while (nCopied < dwLength)
        {
            DWORD nChunkSize= std::min(MEM_PAGE_SIZE-(addr&(MEM_PAGE_SIZE-1)), dwLength-nCopied);

            if (IsBadWritePtr((void*)addr, nChunkSize)) {
                if (!bIgnoreErrors)
                    return false;
            }
            else {
                memcpy((void*)addr, &buf[nCopied], nChunkSize);
            }
            nCopied += nChunkSize;
            addr += nChunkSize;
        }
    }
    else if (IsBadWritePtr((void*)addr, dwLength)) {
        // no direct access.
    }
    else if ((nDataAccess&7)==1) {
        while (nCopied < dwLength)
        {
            *(BYTE*)addr= *(BYTE*)(buf+nCopied);
            nCopied++;
            addr++;
        }
    }
    else if (addr&1) {
        // non-word aligned address not allowed
    }
    else if ((nDataAccess&7)==2) {
        while (nCopied+1 < dwLength)
        {
            *(WORD*)addr= *(WORD*)(buf+nCopied);
            nCopied+=2;
            addr+=2;
        }
    }
    else if (addr&3) {
        // non-dword aligned address not allowed
    }
    else if ((nDataAccess&7)==4) {
        while (nCopied+3 < dwLength)
        {
            *(DWORD*)addr= *(DWORD*)(buf+nCopied);
            nCopied+=4;
            addr+=4;
        }
    }
    *pdwCopied= nCopied;

    return true;
}
bool ReadPhysicalMemory(DWORD dwStart, BYTE *buf, DWORD dwLength, DWORD *pdwCopied, int nDataAccess)
{
    DWORD nCopied= 0;

    DWORD dwRoundedStart= dwStart&KInfoTable[KINX_PFN_MASK];
    DWORD dwStartOffset= dwStart-dwRoundedStart;
    DWORD dwRoundedLength= dwLength+dwStartOffset;

    memdebug("readphysmem(a=%08lx l=%08lx, acc=%d)\n", dwStart, dwLength, nDataAccess);

    BYTE *pVirt= (BYTE*)VirtualAlloc(NULL, dwRoundedLength, MEM_RESERVE, PAGE_NOACCESS);
    if (pVirt==NULL) {
        *pdwCopied= nCopied;
        error("VirtualAlloc(l=%08lx)\n", dwRoundedLength);
        return false;
    }

    if (!VirtualCopy(pVirt, (void*)(dwRoundedStart>>8), dwRoundedLength, PAGE_PHYSICAL|PAGE_READONLY|PAGE_NOCACHE)) {
        error("VirtualCopy(s=%08lx l=%08lx)\n", dwRoundedStart, dwRoundedLength);
        VirtualFree(pVirt, dwRoundedLength, MEM_DECOMMIT);
        VirtualFree(pVirt, dwRoundedLength, MEM_RELEASE);
        *pdwCopied= nCopied;
        return false;
    }

    DWORD addr= (DWORD)pVirt+dwStartOffset;

    bool bRes= ReadMemory(addr, buf, dwLength, pdwCopied, nDataAccess);
    VirtualFree(pVirt, dwRoundedLength, MEM_DECOMMIT);
    VirtualFree(pVirt, dwRoundedLength, MEM_RELEASE);

    return bRes;
}
bool WritePhysicalMemory(DWORD dwStart, const BYTE *buf, DWORD dwLength, DWORD *pdwCopied, DWORD nDataAccess)
{
    DWORD nCopied= 0;

    DWORD dwRoundedStart= dwStart&KInfoTable[KINX_PFN_MASK];
    DWORD dwStartOffset= dwStart-dwRoundedStart;
    DWORD dwRoundedLength= dwLength+dwStartOffset;

    BYTE *pVirt= (BYTE*)VirtualAlloc(NULL, dwRoundedLength, MEM_RESERVE, PAGE_NOACCESS);
    if (pVirt==NULL) {
        *pdwCopied= nCopied;
        return false;
    }

    if (!VirtualCopy(pVirt, (void*)(dwRoundedStart>>8), dwRoundedLength, PAGE_PHYSICAL|PAGE_READWRITE|PAGE_NOCACHE)) {
        VirtualFree(pVirt, dwRoundedLength, MEM_DECOMMIT);
        VirtualFree(pVirt, dwRoundedLength, MEM_RELEASE);
        *pdwCopied= nCopied;
        return false;
    }

    DWORD addr= (DWORD)pVirt+dwStartOffset;

    bool bRes= WriteMemory(addr, buf, dwLength, pdwCopied, nDataAccess);

    VirtualFree(pVirt, dwRoundedLength, MEM_DECOMMIT);
    VirtualFree(pVirt, dwRoundedLength, MEM_RELEASE);

    return bRes;
}
ITSUTILS_API HRESULT STDAPICALLTYPE ITReadProcessMemory(
        DWORD cbInput, ReadProcessMemoryParams *pbInput,
        DWORD *pcbOutput, ReadProcessMemoryResult **ppbOutput,
        IRAPIStream *pStream)
{
    ReadProcessMemoryResult *pOut=NULL;

    *pcbOutput= PTR_DIFF(pOut, pOut->buffer+pbInput->nSize);

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

    KernelMode _km;

    //debug("itrm(h=%08lx)\n", pbInput->hProcess);
    DWORD res= 0;
    if (pbInput->hProcess==INVALID_HANDLE_VALUE)    // -p
        ReadPhysicalMemory(pbInput->dwOffset, &pOut->buffer[0], pbInput->nSize, &pOut->dwNumberOfBytesRead, pbInput->nDataAccess);
    else if (pbInput->hProcess==NULL)   // -m
        ReadMemory(pbInput->dwOffset, &pOut->buffer[0], pbInput->nSize, &pOut->dwNumberOfBytesRead, pbInput->nDataAccess);
    else if (!SafeReadProcMem(pbInput->hProcess, pbInput->dwOffset, &pOut->buffer[0], pbInput->nSize, &pOut->dwNumberOfBytesRead, (pbInput->nDataAccess&8)!=0))
        res= GetLastError();


    return res;
}

ITSUTILS_API HRESULT STDAPICALLTYPE ITWriteProcessMemory(
        DWORD cbInput, WriteProcessMemoryParams *pbInput,
        DWORD *pcbOutput, WriteProcessMemoryResult **ppbOutput,
        IRAPIStream *pStream)
{
    *pcbOutput= sizeof(WriteProcessMemoryResult);
    WriteProcessMemoryResult *pOut= *ppbOutput= (WriteProcessMemoryResult*)LocalAlloc(LPTR, *pcbOutput);

    DWORD res= 0;
    KernelMode _km;

    if (pbInput->hProcess==INVALID_HANDLE_VALUE)
        WritePhysicalMemory(pbInput->dwOffset, &pbInput->buffer[0], pbInput->nSize, &pOut->dwNumberOfBytesWritten, pbInput->nDataAccess);
    else if (pbInput->hProcess==NULL)
        WriteMemory(pbInput->dwOffset, &pbInput->buffer[0], pbInput->nSize, &pOut->dwNumberOfBytesWritten, pbInput->nDataAccess);
    else if (!WriteProcessMemory(pbInput->hProcess, (void*)pbInput->dwOffset, &pbInput->buffer[0], pbInput->nSize, &pOut->dwNumberOfBytesWritten))
        res= GetLastError();


    return res;
}