Get mutants and events from processes

No more words, just codes.

#include "StdAfx.h"
#include "Informer.h"
#include <Tlhelp32.h>

namespace {
    PZwQuerySystemInformation ZwQuerySystemInformation = NULL;
    PZwDuplicateObject ZwDuplicateObject = NULL;
    PZwQueryObject ZwQueryObject = NULL;
    pFnRtlAdjustPrivilege RtlAdjustPrivilege = NULL;
}

CInformer::CInformer(void)
{
    if(!(ZwQuerySystemInformation = (PZwQuerySystemInformation)GetProcAddress(GetModuleHandle(_T("ntdll.dll")), "ZwQuerySystemInformation"))) {
        ASSERT(FALSE);
    }
    if(!(ZwDuplicateObject = (PZwDuplicateObject)GetProcAddress(GetModuleHandle(_T("ntdll.dll")), "ZwDuplicateObject"))) {
        ASSERT(FALSE);
    }
    if(!(ZwQueryObject = (PZwQueryObject)GetProcAddress(GetModuleHandle(_T("ntdll.dll")), "ZwQueryObject"))) {
        ASSERT(FALSE);
    }

    BOOL bRet = FALSE;
    HANDLE hToken;
    bRet = OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &hToken);
    if(!bRet) {
        ASSERT(FALSE);
    }

    bRet = setPrivilege(hToken, SE_DEBUG_NAME, TRUE);

    if(!bRet) {
        ASSERT(FALSE);
    }
}

CInformer::~CInformer(void)
{
}

BOOL CInformer::setPrivilege(HANDLE hToken, LPCTSTR szPrivilege, BOOL bOk) {
    BOOL bRet = FALSE;
    LUID luid;
    TOKEN_PRIVILEGES tp;

    bRet = LookupPrivilegeValue(NULL, szPrivilege, &luid);
    if(!bRet) {
        return bRet;
    }

    tp.PrivilegeCount = 1;

    if(bOk) {
        tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
    } else {
        tp.Privileges[0].Attributes = NULL;
    }

    bRet = AdjustTokenPrivileges(hToken, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), (PTOKEN_PRIVILEGES)NULL, (PDWORD)NULL);

    if(!bRet) {
        return bRet;
    }
    return TRUE;
}

wstring CInformer::GetProcessNameByID(DWORD pid) {
    HANDLE hSnapshot;
    PROCESSENTRY32 pe;
    BOOL fFound = FALSE;

    pe.dwSize = sizeof(PROCESSENTRY32);
    hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if(hSnapshot != NULL) {
        fFound = Process32First(hSnapshot, &pe);
        while(fFound) {
            if(pid == pe.th32ProcessID) {
                CloseHandle(hSnapshot);
                return pe.szExeFile;
            }
            fFound = Process32Next(hSnapshot, &pe);
        }
        CloseHandle(hSnapshot);
    }
    return L"";
}

PULONG GetSystemHandleInformationSize(ULONG n) {
    ULONG nSize = 0x1000;
    PULONG pBuffer;

    pBuffer = new ULONG[nSize];
    NTSTATUS status = ZwQuerySystemInformation(SystemHandleInformation,(PVOID)pBuffer,0x1000,&nSize);
    while(status == STATUS_INFO_LENGTH_MISMATCH) {
        delete [] pBuffer;
        if(STATUS_INFO_LENGTH_MISMATCH == status) {
            pBuffer = new ULONG[nSize];
            status = ZwQuerySystemInformation(SystemHandleInformation,(PVOID)pBuffer,nSize,&nSize);
            if(status == STATUS_INFO_LENGTH_MISMATCH)
                continue;
            n = nSize;
            return pBuffer;
        }
    }
    return 0;
}

void CInformer::GetProcessInformation(vector<ProcessInformation>& vecPI) {
    ULONG n = 0x4000;
    PULONG p = GetSystemHandleInformationSize(n);

    // indicate the array to a handle information structure
    PSYSTEM_HANDLE_INFORMATION h = PSYSTEM_HANDLE_INFORMATION(p + 1);

    // Query handle information

    for(ULONG i = 0; i < * p; i++) {
        HANDLE hProcess;
        // if the process is system process, ingore it.
        if(h[i].ProcessId == 0 || h[i].ProcessId == 4 || (hProcess = OpenProcess(PROCESS_DUP_HANDLE, FALSE, h[i].ProcessId)) == NULL) continue;

        DWORD dwQueryPID = h[i].ProcessId;

        HANDLE hObject;
        if(ZwDuplicateObject(hProcess, HANDLE(h[i].Handle), NtCurrentProcess(), &hObject, 0, 0, DUPLICATE_SAME_ATTRIBUTES) != STATUS_SUCCESS) {
            continue;
        }
        OBJECT_BASIC_INFORMATION obi;
        ZwQueryObject(hObject, ObjectBasicInformation, &obi, sizeof obi, &n);

        n = obi.TypeInformationLength + 2; // extend the length to get the type name

        POBJECT_TYPE_INFORMATION oti = (POBJECT_TYPE_INFORMATION) (new char[n]);

        ZwQueryObject(hObject, ObjectTypeInformation, oti, n, &n);
        n = obi.NameInformationLength == 0 ? MAX_PATH* sizeof(WCHAR): obi.NameInformationLength;

        if(wcscmp(oti[0].TypeName.Buffer, L"File") == 0) {
            CloseHandle(hObject);
            delete [] oti;
            continue;
        }

        POBJECT_NAME_INFORMATION oni = (POBJECT_NAME_INFORMATION) (new char[n]);

        NTSTATUS rv = ZwQueryObject(hObject, ObjectNameInformation, oni, n, &n);

        if(STATUS_SUCCESS != rv) {
            delete [] (char*)oti;
            delete [] (char*)oni;
            CloseHandle(hObject);
            continue;
        }


        if(oni[0].Name.Buffer == NULL || wcscmp(oni[0].Name.Buffer, L"") == 0) {
            delete [] (char*)oti;
            delete [] (char*)oni;
            CloseHandle(hObject);
            continue;
        }

        ProcessInformation pi;

        if(wcscmp(oti[0].TypeName.Buffer, L"Event") != 0 && wcscmp(oti[0].TypeName.Buffer, L"Mutant") != 0) {
            goto clean;
        }

        if(wcscmp(oti[0].TypeName.Buffer, L"") == 0)
            goto clean;

        if(wcscmp(oti[0].TypeName.Buffer, L"Event") == 0) {
            pi.EventName = oni[0].Name.Buffer;
        } else if(wcscmp(oti[0].TypeName.Buffer, L"Mutant") == 0) {
            pi.MutexName = oni[0].Name.Buffer;
        }
        pi.handle = HANDLE(h[i].Handle);
        pi.dwPID = h[i].ProcessId;
        pi.strProcessName = GetProcessNameByID(h[i].ProcessId);
        vecPI.push_back(pi);

clean:
        delete [] (char*)oti;
        delete [] (char*)oni;
        CloseHandle(hObject);

    }
    delete [] p;
}