\
#include "pch.h"
#include "process_notify.h"
#include "psapi.h"
#include "vector"
#include "algorithm"
#include "strsafe.h"
#pragma comment(lib, "Psapi.lib")


CProcessNotify::CProcessNotify()
    :
    m_hThread(NULL),
    m_hQuit(NULL)
{
}

CProcessNotify::~CProcessNotify()
{
    FinishPrivate();
}

BOOL CProcessNotify::Start(DWORD nEnumInterval, BOOL bNotifyCurrent, fnNotifyCallBack cb, BOOL bOnce)
{
    if (m_hThread)
        return TRUE;

    if (!m_hQuit)
    {
        m_hQuit = CreateEvent(NULL, TRUE, FALSE, NULL);
        if (!m_hQuit)
            return FALSE;
    }

    m_bOnce = bOnce;
    m_nEnumInterval = nEnumInterval;
    m_bNotifyCurrent = bNotifyCurrent;
    m_cb = cb;

    DWORD nThreadId = 0;
    m_hThread = CreateThread(NULL, 0, ProcessNotifyProc, (PVOID)this, 0, &nThreadId);
    if (!m_hThread)
        return FALSE;

    return TRUE;
}

DWORD
CProcessNotify::StartPrivate()
{
    typedef std::vector<DWORD> PIDLIST;
    typedef std::vector<DWORD>::iterator PIDLISTITER;

    const DWORD nSize = 1024;
    PIDLIST vPidListPrev;
    DWORD dwInterval = 1;

    for (;;)
    {
        // ÃÖÃÊ´Â ¹Ù·Î ½ÃÀÛ
        DWORD nRet = WaitForSingleObject(m_hQuit, dwInterval);
        if (WAIT_TIMEOUT != nRet)
            return nRet;

        dwInterval = m_nEnumInterval;
        PIDLIST vPidListNext(nSize);
        DWORD nNeeded = 0;
        BOOL bRet = EnumProcesses(&vPidListNext[0], nSize * sizeof(DWORD), &nNeeded);
        if (!bRet)
            return GetLastError();

        vPidListNext.resize(nNeeded / sizeof(DWORD));
        sort(vPidListNext.begin(), vPidListNext.end());

        if (vPidListPrev.empty())
        {
            if (!m_bNotifyCurrent)
            {
                vPidListPrev = vPidListNext;
                continue;
            }
        }

        if (vPidListNext.size() == vPidListPrev.size())
        {
            int nRet = memcmp(&vPidListNext[0], &vPidListPrev[0], nNeeded);
            if (!nRet)
                continue;
        }

        PIDLISTITER it;
        PIDLIST vPidListAdded(nSize);
        PIDLISTITER itAdded = set_difference(vPidListNext.begin(), vPidListNext.end(), vPidListPrev.begin(), vPidListPrev.end(), vPidListAdded.begin());
        for (it = vPidListAdded.begin(); it != itAdded; it++)
        {
            BOOL bRet = OnNotify(*it, TRUE);
            if (!bRet)
                return ERROR_CANCELLED;

            DWORD nRet = WaitForSingleObject(m_hQuit, 0);
            if (nRet != WAIT_TIMEOUT)
                return nRet;
        }

        PIDLIST vPidListRemoved(nSize);
        PIDLISTITER itRemoved = set_difference(vPidListPrev.begin(), vPidListPrev.end(), vPidListNext.begin(), vPidListNext.end(), vPidListRemoved.begin());
        for (it = vPidListRemoved.begin(); it != itRemoved; it++)
        {
            BOOL bRet = OnNotify(*it, FALSE);
            if (!bRet)
                return ERROR_CANCELLED;

            DWORD nRet = WaitForSingleObject(m_hQuit, 0);
            if (nRet != WAIT_TIMEOUT)
                return nRet;
        }

        vPidListPrev = vPidListNext;
        if (m_bOnce)
            break;

    }

    return 0;
}



DWORD WINAPI CProcessNotify::ProcessNotifyProc(LPVOID pArg)
{
    CProcessNotify* pProcessNotify = reinterpret_cast<CProcessNotify*>(pArg);
    if (!pProcessNotify)
    {
        return ERROR_INVALID_PARAMETER;
    }
    return pProcessNotify->StartPrivate();
}

BOOL
CProcessNotify::Finish()
{
    return FinishPrivate();
}

BOOL
CProcessNotify::FinishPrivate()
{
    if (!m_hThread)
    {
        SetLastError(ERROR_INVALID_PARAMETER);
        return FALSE;
    }

    if (m_hQuit)
    {
        SetEvent(m_hQuit);
    }

    WaitForSingleObject(m_hThread, INFINITE);
    CloseHandle(m_hThread);
    CloseHandle(m_hQuit);
    m_hThread = NULL;
    m_hQuit = NULL;

    return TRUE;
}



BOOL GetProcessFullPath(DWORD dwPid, LPWSTR lpwProcessFullPath, DWORD dwSize)
{
    static ZWQUERYINFORMATIONPROCESS query = NULL;

    NTSTATUS dwError = 0;
    HANDLE hProcess = NULL;
    DWORD dwReturnSize = 0;
    DWORD bufsize = dwReturnSize;
    PVOID* buf = NULL;
    PUNICODE_STRING pUni = NULL;

    hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, dwPid);
    if (hProcess == NULL)
    {
        return FALSE;
    }

    if (query == NULL)
    {
        query = (ZWQUERYINFORMATIONPROCESS)GetProcAddress(GetModuleHandleA("ntdll.dll"), "NtQueryInformationProcess");
    }

    dwError = query(hProcess, ProcessImageFileName, NULL, 0, &dwReturnSize);
    if (dwError != STATUS_INFO_LENGTH_MISMATCH)
    {
        goto $GetProcessFullPathCleanup;
    }

    
    buf = (PVOID*)malloc(bufsize);
    if (!buf)
        goto $GetProcessFullPathCleanup;

    memset(buf, 0, bufsize);
    dwError = query(hProcess, ProcessImageFileName, buf, bufsize, &dwReturnSize);
    if (dwError != 0)
    {
        goto $GetProcessFullPathCleanup;
    }

    pUni = (PUNICODE_STRING)buf;

    //OutputDebugStringW(pUni->Buffer);
    StringCbCopyNW(lpwProcessFullPath, dwSize, pUni->Buffer, pUni->Length);

$GetProcessFullPathCleanup:
    if (hProcess)
        CloseHandle(hProcess);

    if (buf)
        free(buf);


    return TRUE;
}

BOOL
CProcessNotify::OnNotify(DWORD nPid, BOOL bCreate)
{
    if (bCreate)
    {
        WCHAR FileName[1024] = { 0, };
        BOOL i = GetProcessFullPath(nPid, FileName, sizeof(FileName));
        m_cb(bCreate, nPid, FileName, (DWORD)wcslen(FileName));
    }
    else
    {
        m_cb(bCreate, nPid, NULL, 0);
    }

    return TRUE;
}