#include <windows.h>
#include <debug.h>
#include <algorithm>
#include "itsutils.h"
#include "kernelmisc.h"
#include "ptrutils.h"
#include "stringutils.h"
#if 0
// todo: create layout like this:
// FNOT  <hproc>  <pvobj>
// <pvobj>: <dword hexdump>
// {
//   <ptr1>: <dword hexdump>
//   <ptr2>: <dword hexdump>
//   {
//     ptr2ptr1: <dword hexdump>
//   }
//   <ptr3>: <dword hexdump>
// }
//
struct ptrinfo {
    ptrinfo() : offset(0), level(0) { }
    DWORD offset;
    int level;
    DwordVector data;
};
typedef std::map<DWORD,ptrinfo> ptrinfo_map;
typedef std::set<DWORD> usedptrs;
void getstruct(PROCESS *objproc, DWORD dwOfs, int level)
{
    ptrinfo_map ptrmap;
    DwordVector v;
    if (!GetDwordData(objproc, dwOfs, v))
        return;

    typedef std::pair<ptrinfo_map::iterator,bool> insertresult;
    insertresult i= ptrmap.insert(ptrinfo_map::value_type(dwOfs, ptrinfo));
    if ((*i).second) {
        (*(*i).first).level= level;
        GetDwordData(objproc, dwOfs, (*(*i).first).data);
    }
}
bool GetDwordData(PROCESS *objproc, DWORD dwOfs, DwordVector& v)
{
    if (dwOfs&3) return false;
    if (IsBadReadPtr(dwOfs-8, 0x18)) return false;
    if (dwOfs==NULL || objproc==NULL) return false;
    DWORD head[2];
    DWORD nRead;
    if (!ReadProcessMemory(objproc->hProc,  (void*)(dwOfs-8), head, sizeof(head), &nRead)) return false;
    if (IsBadReadPtr((void*)head[1],0x40)) return false;
    if (head[0]>0x1000) return false; 
    DWORD tail[2];
    if (!ReadProcessMemory(objproc->hProc,  (void*)(dwOfs-8+head[0]), tail, sizeof(tail), &nRead)) return false;
    if (head[1]!=tail[1]) return false;

    v.resize((head[0]-8)/4);
    if (!ReadProcessMemory(objproc->hProc,  (void*)dwOfs, &v[0], v.size()*sizeof(DWORD), &nRead)) return false;

    return true;
}
#endif
std::string processDescription(PROCESS *p)
{
    return stringformat("%08lx:%08lx:'%ls'", p->hProc, p->dwVMBase, p->lpszProcName);
}

std::string processDescription(HANDLE hProc)
{
    HDATA *hd= cvHandle2HDataPtr(hProc);
    return processDescription((PROCESS*)hd->pvObj);
}

std::string moduleDescription(Module *m)
{
    return stringformat("%08lx:'%ls'", m, m->lpszModName);
}

std::string getHandleDescription(const std::string& type, PROCESS *objproc, void *pvObj)
{
    if (pvObj==NULL)
        return "(null)";
    if (type=="PROC") {
        PROCESS *p= (PROCESS*)pvObj;
        return stringformat("%s %ls", processDescription(p).c_str(), p->pcmdline);
    }
    else if (type=="THRD") {
        THREAD *t= (THREAD*)pvObj;
        return stringformat("%s start=%08lx curip=%08lx sp=%08lx", processDescription(t->pOwnerProc).c_str(),
                t->dwStartAddr, t->ctx.Pc, t->ctx.Sp);
    }
    else if (type=="APIS") {
        APISET *a= (APISET*)pvObj;
        return stringformat("%-4s reg=%02x dsp=%02x type=%02x #%d, pfn=%08x %s", std::string(a->cinfo.acName, a->cinfo.acName+4).c_str(),
                a->iReg, a->cinfo.disp, a->cinfo.type, a->cinfo.cMethods, a->cinfo.ppfnMethods,
                processDescription(a->cinfo.pServer).c_str());
    }
    else if (type=="EVNT") {
        EVENT *e= (EVENT*)pvObj;
        return stringformat("%ls", e->name?e->name->name:L"");
    }
    else if (type=="MUTX") {
        MUTEX *m= (MUTEX*)pvObj;
        return stringformat("%d %ls", m->LockCount, m->name?m->name->name:L"");
    }
    else if (type=="SEMP") {
        SEMAPHORE *s= (SEMAPHORE*)pvObj;
        return stringformat("%d/%d/%d %ls", s->lCount, s->lMaxCount, s->lPending, s->name?s->name->name:L"");
    }
    else if (type=="FMAP") {
        FSMAP*f= (FSMAP*)pvObj;
        return stringformat("f=%08lx,l=%08x regionlen=%08lx  event=%ls", f->hFile, f->filelen, f->length, (f->name?f->name->name:L""));
    }
    else if (type=="W32D") {
        DWORD opendev[7];
        // 0:    struct fsopendev_t *nextptr;  // next one in linked list
        // 1:    DWORD dwOpenData;
        // 2:    fsdev_t *lpDev;
        // 3:    DWORD *lpdwDevRefCnt;   // since we set lpDev to NULL on device deregister
        // 4:    DWORD dwOpenRefCnt;     // ref count for this structure 
        // 5:    HANDLE KHandle;         // kernel handle pointing to this structure
        // 6:    HPROCESS hProc;         // process owning this handle
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, opendev, sizeof(opendev), &nRead)) {
            return "";
        }
        DWORD devdesc[24];       // opendev[3] = &(lpDev->dwRefCnt)
//  0:-16:  pInitFn fnInit;             // required
//  1:-15:  pPreDeinitFn fnPreDeinit;   // optional
//  2:-14:  pDeinitFn fnDeinit;         // required
//  3:-13:  pOpenFn fnOpen;             // if present, need close and at least one of read/write/seek/control
//  4:-12:  pPreCloseFn fnPreClose;     // optional even if close is present
//  5:-11:  pCloseFn fnClose;           // required if open is present
//  6:-10:  pReadFn fnRead;
//  7: -9:  pWriteFn fnWrite;
//  8: -8:  pSeekFn fnSeek;
//  9: -7:  pControlFn fnControl;
// 10: -6:  pPowerupFn fnPowerup;       // optional, can be NULL
// 11: -5:  pPowerdnFn fnPowerdn;       // optional, can be NULL
// 12: -4:  DWORD dwData;               // instance data returned from Init() entry point
// 13: -3:  HINSTANCE hLib;             // driver dll handle
// 14: -2:  DWORD dwId;                 // active key ID
// 15: -1:  WORD wFlags;                // DF_XXX state flags
// 16:  0:  DWORD dwRefCnt;             // use count for this device structure
// 17:  1:  HANDLE hParent;             // parent device driver, or NULL
// 18:  2:  fsinterface_t *pInterfaces; // interfaces advertised by this device instance
// 19:  3:  LPWSTR pszLegacyName;       // e.g., "COM1:"
// 20:  4:  LPWSTR pszDeviceName;       // e.g., "Serial1" in $device namespace
// 21:  5:  LPWSTR pszBusName;          // e.g., "Builtin_1_1" in $bus namespace
// 22:  6:  LPWSTR pszDeviceKey;        // NULL if not an activated device
// 23:  7:  DWORD dwActivateId;         // thread ID for the thread calling ActivateDevice() (only valid during Init())
        if (!ReadProcessMemory(objproc->hProc,  (void*)(opendev[3]-16*4), devdesc, sizeof(devdesc), &nRead)) {
            return "";
        }
        WCHAR devname[8];       // devdesc[19] = pszLegacyName
        if (!ReadProcessMemory(objproc->hProc,  (void*)devdesc[19], devname, sizeof(devname), &nRead)) {
            return "";
        }
        return stringformat("%-8ls data=%08lx %s  %s", devname, 
                opendev[1],     // dwData
                processDescription((HANDLE)opendev[6]).c_str(),
                moduleDescription((Module*)devdesc[13]).c_str());
    }
    else if (type=="HFSD" ||  type=="FFSD") {
        struct hfsdinfo {
            DWORD nexthfsd;
            DWORD prevhfsd;
            DWORD partobj;
            DWORD w3;
            DWORD pfileinfo;  // for ffsd
            DWORD hOwnerProc;
            DWORD hFile;
            DWORD pgtgt;
        };
        hfsdinfo hi;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &hi, sizeof(hi), &nRead)) {
            return "";
        }
        if (type=="FFSD") {
            return stringformat("%s", processDescription((HANDLE)hi.hOwnerProc).c_str());
        }
        struct gtgtinfo {
            DWORD magic;
            DWORD w1;
            DWORD pdgdg;  // parent directory
            DWORD pname;
            DWORD prevgtgt;
            DWORD nextgtgt;
        };
        gtgtinfo gt;
        if (!ReadProcessMemory(objproc->hProc,  (void*)hi.pgtgt, &gt, sizeof(gt), &nRead)) {
            return "";
        }
        struct dgdginfo {
            DWORD magic;
            DWORD pname;
            DWORD ptgtg;
            DWORD pgdgd;
            DWORD pfirstgtgt;
            DWORD pfirstchild;
            DWORD parent;
            DWORD nextsib;
            DWORD prevsib;
        };
        // todo: process dgdg, to get full path.
        WCHAR filename[MAX_PATH];
        if (!ReadProcessMemory(objproc->hProc,  (void*)gt.pname, filename, sizeof(filename), &nRead)) {
            return "";
        }
        return stringformat("%s  %ls", processDescription((HANDLE)hi.hOwnerProc).c_str(), filename);
    }
    else if (type=="FNOT") {  // 237 handles
        // registry notifications
        struct gdgdinfo {
            DWORD magic;
            DWORD hProc;
            DWORD w2,w3;
            DWORD hEvent;
            DWORD hNotify;
            DWORD w6;
            DWORD ptgtg;   // -> HIVE
            DWORD w8;
            DWORD w9;
            DWORD pdgdg;   // -> key
            DWORD prev;
            DWORD next;
        };
        gdgdinfo gi;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &gi, sizeof(gi), &nRead)) {
            return "";
        }

        struct tgtginfo {
            DWORD magic;
            DWORD w1;
            DWORD w2;
            DWORD w3;
            DWORD w4;
            DWORD w5;
            DWORD hivename;
            DWORD w7;
            DWORD pdgdg;   // -> first key
            DWORD pgdgd;   // -> first notification
            DWORD w10;
            DWORD next;
        };
        tgtginfo ti;
        if (!ReadProcessMemory(objproc->hProc,  (void*)gi.ptgtg, &ti, sizeof(gi), &nRead)) {
            return "";
        }
        TCHAR hivename[MAX_PATH];
        if (!ReadProcessMemory(objproc->hProc,  (void*)ti.hivename, hivename, sizeof(hivename), &nRead)) {
            return "";
        }
        

        struct dgdginfo {
            DWORD magic;
            DWORD pname;        // keyname
            DWORD ptgtg;        // -> hive
            DWORD pgdgd;
            DWORD pfirstgtgt;
            DWORD pfirstchild;
            DWORD parent;
            DWORD nextsib;
            DWORD prevsib;
        };
        WStringList path;

        DWORD p= gi.pdgdg;
        while (p) {
            dgdginfo dgi;
            if (!ReadProcessMemory(objproc->hProc,  (void*)p, &dgi, sizeof(dgi), &nRead)) {
                return "";
            }

            path.resize(path.size()+1);
            path.back().resize(MAX_PATH);

            if (!ReadProcessMemory(objproc->hProc,  (void*)dgi.pname, &path.back()[0], MAX_PATH, &nRead)) {
                return "";
            }
            path.back().resize(_tcslen(path.back().c_str()));

            p= dgi.parent;
        }
        std::reverse(path.begin(), path.end());

        return stringformat("GDGD:%08x:%s TGTG[%08x]:%ls DGDG[%08x]:%ls",
                gi.hNotify, processDescription((HANDLE)gi.hProc).c_str(),
                gi.ptgtg, hivename,
                gi.pdgdg, JoinStringList(path, _T("\\")).c_str());
    }
    else if (type=="EDBO") {  // 30 handles
        struct edboinfo {
            DWORD w1;
            DWORD w12345678;  // magic
            DWORD n;
            DWORD p1;
            DWORD nul1;
            DWORD nul2;
            DWORD plock;
            DWORD nul3;
            DWORD nul4;
            DWORD pvolume;
            GUID  guid;
        };
        struct edbovolume {     // size = 0xf8
            DWORD w1;
            DWORD hMutx;
            DWORD w2;
            DWORD w3;
            DWORD p1;
            DWORD p2;
            DWORD w4;
            DWORD w5;
            DWORD pvolumehandles;
        };
        struct edbovolhandles {
            DWORD nul1;
            DWORD hProc;
            DWORD hHFSD;
            DWORD hMUTX1;
            DWORD hMUTX2;
            DWORD hFMAP;
            DWORD hMutx3;
            WCHAR volumename[MAX_PATH];
        };
        // todo
    }
    else if (type=="DBOA") {  // 28 handles
        struct dboainfo {
            DWORD next;
            DWORD hdboa;
            DWORD hproc;
            DWORD dw3;
            DWORD pvolume;
            DWORD oid;
            DWORD dw6;
            DWORD dw7;
        };
        dboainfo di;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &di, sizeof(di), &nRead)) {
            return "";
        }
        struct dbvolinfo {
            DWORD next;
            DWORD dw1;
            DWORD pfilebuffer;
            DWORD dw3;
            TCHAR filename[MAX_PATH];
            // ... more
        };
        dbvolinfo vi;
        if (!ReadProcessMemory(objproc->hProc,  (void*)di.pvolume, &vi, sizeof(vi), &nRead)) {
            return "";
        }
        return stringformat("%s oid=xx%06x %ls", processDescription((HANDLE)di.hproc).c_str(), di.oid, vi.filename);
    }
    else if (type=="Sock") {  // 21 handles
        // 1 dword containing the handle value.
    }
    else if (type=="W32S") {  // 13 handles
        struct w32sinfo {
            DWORD dw0, dw1;
            DWORD pdrvrinfo;
            DWORD dw3;
            DWORD hSvc;
            DWORD hProc;
        };
        w32sinfo si;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &si, sizeof(si), &nRead)) {
            return "";
        }
        struct drvrinfo {
            DWORD nextdrvr;
            DWORD prevdrvr;
            DWORD dw[11];
            TCHAR name[3];
            WORD  index;
            DWORD w15;
            DWORD pdllname;
            DWORD w17;
            DWORD pregkey;
            DWORD w19;
        };
        drvrinfo di;
        if (!ReadProcessMemory(objproc->hProc,  (void*)si.pdrvrinfo, &di, sizeof(di), &nRead)) {
            return "";
        }
        TCHAR dllname[MAX_PATH];
        if (!ReadProcessMemory(objproc->hProc,  (void*)di.pdllname, dllname, sizeof(dllname), &nRead)) {
            return "";
        }
        TCHAR regkey[MAX_PATH];
        if (!ReadProcessMemory(objproc->hProc,  (void*)di.pregkey, regkey, sizeof(regkey), &nRead)) {
            return "";
        }
        return stringformat("%s  %c%c%c%d:  %ls  %ls",
                processDescription((HANDLE)si.hProc).c_str(), di.name[0], di.name[1], di.name[2], di.index, dllname, regkey);
    }
    else if (type=="STRG") {  // 4 handles
        struct strginfo {
            DWORD magic;
            DWORD p1;   // driver info
            DWORD p2;   // partname
            DWORD w3;
            DWORD w4;
            DWORD pnext;
        };
        // todo: dump driver + partition name
    }
    else if (type=="PFSD") {
        struct partinfo {
            DWORD w0;  // prevpfsd
            DWORD w1;  // nextpfsd
            DWORD w2;  // fileobj  hfsd.0 points here
            DWORD w3;  // fileobj
            DWORD w4;
            DWORD w5;  // -> tgtg obj
            DWORD w6;
            DWORD w7;
            DWORD w8;
            DWORD w9;
            DWORD w10[0x21];
            DWORD w43[4];
            WCHAR devname[32];
            WCHAR partname[32];
        };
        partinfo pi;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &pi, sizeof(pi), &nRead)) {
            return "";
        }
        return stringformat("%ls:%ls", pi.devname, pi.partname);
    }
    else if (type=="DMFS") {  // 2 handles
        // no object ptr
    }
    else if (type=="ROM3") {  // 1 handle
        // no object ptr
    }
    else if (type=="ROM1") {
        // 00000000 00000000 80000021 00000000 00000000 00000000
    }
    else if (type=="PRdr") {
        // deadface 00000000
    }
    else if (type=="FFHN") {
        struct ffhninfo {
            DWORD dw[9];
            WORD  len;
            TCHAR pattern[MAX_PATH];
        };
        ffhninfo ff;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &ff, sizeof(ff), &nRead)) {
            return "";
        }
        return ToString(std::Wstring(ff.pattern, &ff.pattern[ff.len]));
    }
    else if (type=="DMFF") {
        // contains ptr to devmgr.dll
    }
    else if (type=="BDEV") {
        struct bdevinfo {
            DWORD dw[11];
            TCHAR regkey[MAX_PATH];
        };
        bdevinfo di;
        DWORD nRead;
        if (!ReadProcessMemory(objproc->hProc,  (void*)pvObj, &di, sizeof(di), &nRead)) {
            return "";
        }
        return stringformat("%ls", di.regkey);
    }
    return "";
}
ITSUTILS_API HRESULT STDAPICALLTYPE ITGetHandles(
        DWORD cbInput, GetHandlesParams *pbInput,
        DWORD *pcbOutput, GetHandlesResult **ppbOutput,
        IRAPIStream *pStream)
{
    GetHandlesResult *pOut=NULL;

    KernelMode _km;

    HandleInfoVector v;

    //...

    HDATA *hi= cvHandle2HDataPtr((HANDLE)GetCurrentThreadId());
    HDATA *ha;
    ha=hi;

    unsigned i=0;
    unsigned total=0;
    std::string strbuf;
    v.resize(1024);
    for ( ; total==0 || (ha!=hi && ha!=NULL) ; ha= (HDATA*)ha->linkage.fwd) {
        if ((DWORD(ha->hValue)&3)!=2)
            continue;
        total++;
//      debug("ha=%08lx ", ha);
//      debug("ha.h=%08lx ha.pci=%08lx ", ha->hValue, ha->pci);
        char type[8];
        memcpy(type, ha->pci->acName, 4); memset(type+4, 0, 4);
//      debug("ha.pci.type=%s ha.pci.svr=%08lx ", type, ha->pci->pServer);
//      debug("hProc=%08lx\n", ha->pci->pServer ? ha->pci->pServer->hProc : 0);


        if (pbInput->hProc==INVALID_HANDLE_VALUE || ha->pci->pServer==NULL || ha->pci->pServer->hProc==pbInput->hProc)
            if (pbInput->type[0]==0 || strstr(pbInput->type, type))
            {
                if (i>=v.size())
                    v.resize(v.size()*2);
                HandleInfo &hi= v[i++];

                hi.handle= ha->hValue;
                memcpy(hi.type, type, 8);
                hi.hProc= ha->pci->pServer ? ha->pci->pServer->hProc : NULL;
                hi.pvObj= (DWORD)ha->pvObj;
                hi.descstart= strbuf.size();
                strbuf += getHandleDescription(type, ha->pci->pServer, ha->pvObj);
                hi.descend= strbuf.size();

                //dumpstruct(ha->pci->pServer, ha->pvObj, 0);
            }
    }
    debug("end: ha=%08lx hi=%08lx  count=%d\n", ha, hi, total);

    v.resize(i);
    //...
    *pcbOutput= PTR_DIFF(pOut, pOut->data+v.size())+strbuf.size();
    pOut= *ppbOutput= (GetHandlesResult*)LocalAlloc(LPTR, *pcbOutput);
    pOut->count= v.size();
    memcpy(pOut->data, &v[0], v.size()*sizeof(HandleInfo));
    memcpy(pOut->data+v.size(), strbuf.c_str(), strbuf.size());

    return 0;
}