#include <map>
#include <string>
#include <locale>
#include <windows.h>
#include <winnt.h>

#pragma warning(disable:  4127 ) // conditional expression is constant


using namespace std;

//---------------------------------------------------------------------------------------------------------------------
// Some definitions from NTDDK

typedef LONG	NTSTATUS;
typedef LONG	KPRIORITY;

#define NT_SUCCESS(Status) ((NTSTATUS)(Status) >=  0 )

#define STATUS_INFO_LENGTH_MISMATCH      ((NTSTATUS)0xC0000004L)

#define SystemProcessesAndThreadsInformation	 5 

typedef struct _IDCLIENT 
{
    DWORD	    UniqueProcess;
    DWORD	    UniqueThread;
} IDCLIENT;

typedef struct _UNISTRING 
{
    USHORT	    Length;
    USHORT	    MaximumLength;
    PWSTR	    Buffer;
} UNISTRING;

typedef struct _VM_COUNTERS 
{
    SIZE_T	    PeakVirtualSize;
    SIZE_T	    VirtualSize;
    ULONG	    PageFaultCount;
    SIZE_T	    PeakWorkingSetSize;
    SIZE_T	    WorkingSetSize;
    SIZE_T	    QuotaPeakPagedPoolUsage;
    SIZE_T	    QuotaPagedPoolUsage;
    SIZE_T	    QuotaPeakNonPagedPoolUsage;
    SIZE_T	    QuotaNonPagedPoolUsage;
    SIZE_T	    PagefileUsage;
    SIZE_T	    PeakPagefileUsage;
} VM_COUNTERS;

typedef struct _SYS_THREADS 
{
    LARGE_INTEGER   KernelTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   CreateTime;
    ULONG			WaitTime;
    PVOID			StartAddress;
    IDCLIENT	    ClientId;
    KPRIORITY	    Priority;
    KPRIORITY	    BasePriority;
    ULONG			ContextSwitchCount;
    LONG			State;
    LONG			WaitReason;
} SYS_THREADS, * PSYS_THREADS;

// Note that the size of the SYS_PROCESSES structure is different on
// NT 4 and Win2K, but we don't care about it, since we don't access neither
// IoCounters member nor Threads array

typedef struct _SYS_PROCESSES 
{
    ULONG			NextEntryDelta;
    ULONG			ThreadCount;
    ULONG			Reserved1[ 6 ];
    LARGE_INTEGER   CreateTime;
    LARGE_INTEGER   UserTime;
    LARGE_INTEGER   KernelTime;
    UNISTRING  ProcessName;
    KPRIORITY	    BasePriority;
    ULONG			ProcessId;
    ULONG			InheritedFromProcessId;
    ULONG			HandleCount;
    ULONG			Reserved2[ 2 ];
    VM_COUNTERS	    VmCounters;
#if _WIN32_WINNT >= 0x500
    IO_COUNTERS	    IoCounters;
#endif
    SYS_THREADS  Threads[ 1 ];
} SYS_PROCESSES, * PSYS_PROCESSES;

//---------------------------------------------------------------------------------------------------------------------
//  lstEnumProcesses
//
//  Enumerates all processes in the system using NtQuerySystemInformation.
//  
//  Parameters:
//	[OUT]  lst - list of PID and names
//
//  Returns:
//	  true, if successful, false - otherwise.
//
//---------------------------------------------------------------------------------------------------------------------
// convert wstring_to_string
inline std::string lexical_cast(const std::wstring &source)
    {
        std::string result;
        result.resize( source.size() );

        for ( std::wstring ::const_iterator it = source.begin(); it != source.end(); ++it )
            result[ std::distance ( source.begin(), it )  ] =
                std::use_facet < std::ctype < wchar_t > > ( std::locale(),  0  ).narrow( (*it), '?' );
        return result;
    }
//---------------------------------------------------------------------------------------------------------------------
bool LstEnumProcesses( std::map<std::string, DWORD> &lst )
{
    lst.clear();
	HINSTANCE hpNtDll;
	NTSTATUS (WINAPI * pfnZwQSI)(UINT, PVOID, ULONG, PULONG);
    bool bRes = true;
    PSYS_PROCESSES pProcesses = NULL;
    PVOID pBuffer = NULL;
	NTSTATUS Status;
    ULONG cbBuffer = 0xABCD; // 40K buffer should be enough
	HANDLE hHeap = GetProcessHeap(); // get a handle to the default process heap

	// get handle to NTDLL.DLL
	hpNtDll = GetModuleHandle( "ntdll.dll" );
	if( hpNtDll == NULL )
    {
        return false;
    }
	// find the address of ZwQuerySystemInformation
	*(FARPROC *)&pfnZwQSI = GetProcAddress(hpNtDll, "ZwQuerySystemInformation");
	if( pfnZwQSI == NULL )
    {
        bRes = false;
		goto _end;
    }
    do
    {
		pBuffer = HeapAlloc(hHeap,  0 , cbBuffer);
		if (pBuffer == NULL)
        {
            bRes = false;
		    goto _end;
        }
		Status = pfnZwQSI( SystemProcessesAndThreadsInformation, pBuffer, cbBuffer, NULL);

		if (Status == STATUS_INFO_LENGTH_MISMATCH)
		{
			HeapFree(hHeap,  0 , pBuffer);
			cbBuffer = cbBuffer <<  2 ; // if buffer size is not enough then encrease it
		}
		else if( !NT_SUCCESS(Status) )
		{
			HeapFree(hHeap,  0 , pBuffer);
            bRes = false;
		    goto _end;
		}
    }while( STATUS_INFO_LENGTH_MISMATCH == Status );
    pProcesses = (PSYS_PROCESSES)pBuffer;
    while(true)
    {
		PCWSTR pszProcessName = pProcesses->ProcessName.Buffer;
		if( pszProcessName == NULL )
        {
			pszProcessName = L"System Idle Processes";
        }
#ifdef _UNICODE
        lst[pszProcessName] = pProcesses->ProcessId;
#else
		CHAR szProcessName[MAX_PATH];
		WideCharToMultiByte(CP_ACP,  0 , pszProcessName, - 1 , szProcessName, MAX_PATH, NULL, NULL);

        wstring instr = pszProcessName;
        string outstr = lexical_cast( instr );
        lst[outstr] = pProcesses->ProcessId;
#endif
		if( pProcesses->NextEntryDelta ==  0 )
        {
			break;
        }
		// find the address of the next process structure
		pProcesses = (PSYS_PROCESSES)(((LPBYTE)pProcesses) + pProcesses->NextEntryDelta);
	}
_end:
    if( NULL != pBuffer )
    {
	    HeapFree(hHeap,  0 , pBuffer);
    }
    if(  0  != hpNtDll )
    {
        FreeLibrary(hpNtDll);
    }
	return true;
}
//---------------------------------------------------------------------------------------------------------------------
map<string, string> GetServiceList() // get list of services   <service name, display name>
{
    map<string, string> lstResult;
	SC_HANDLE hSCM = OpenSCManager(NULL, NULL, SC_MANAGER_ENUMERATE_SERVICE);
	if( hSCM == NULL )
    {
		return lstResult;
    }
	BYTE bBuffer[ 2048 ] = { 0 };
	DWORD dwResumeHandle =  0 ;
	DWORD cbNeeded;
	DWORD cServices;
	BOOL bContinue;
	DWORD dwError;
	LPENUM_SERVICE_STATUS pStatus;
	do
	{
		bContinue = FALSE;
		pStatus = (LPENUM_SERVICE_STATUS)bBuffer;
		// retrieve a portion of service status information
		if (!EnumServicesStatus(hSCM, SERVICE_WIN32, SERVICE_STATE_ALL,
								pStatus, sizeof(bBuffer), &cbNeeded,
								&cServices, &dwResumeHandle))
		{
			dwError = GetLastError();
			if (dwError != ERROR_MORE_DATA)
			{
				break;
			}
			bContinue = TRUE;
		}
		for (DWORD i =  0 ; i < cServices; i++) 
		{
            lstResult[ pStatus->lpServiceName ] = pStatus->lpDisplayName;
			pStatus++;
		}
	}
	while( bContinue );
	CloseServiceHandle(hSCM);
    return lstResult;
}