// enumerate processes and notifications of new and ended ones
#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <stdio.h>
#include <ole2.h>
#include <wbemidl.h>

#pragma comment(lib, "wbemuuid.lib")

#define CHECK_HR_AND_RETURN(x) \
	if(FAILED(hr = (x))) \
	{ \
		printf(#x " failed with hr=%x\n", hr); \
		return 1; \
	}

#define CHECK_HR(x) \
	if(FAILED(hr = (x))) \
	{ \
		printf(#x " failed with hr=%x\n", hr); \
	}

#define CONCAT_INT(a, b) a##b
#define CONCAT(a, b) CONCAT_INT(a, b)

#define STACK_BSTR(contents, varName) \
	struct CONCAT(StackBStr, __COUNTER__) \
	{ \
		ULONG bstrSize; \
		WCHAR text[sizeof(contents) / sizeof(WCHAR)]; \
		operator BSTR() {return text;} \
	} varName = {sizeof(contents), contents}

class ProcessNotifier : public IWbemObjectSink
{
	STDMETHOD(QueryInterface)(REFIID iid, PVOID* ppv)
	{
		*ppv = NULL;
		if(iid == IID_IUnknown || iid == IID_IWbemObjectSink)
		{
			AddRef();
			*ppv = this;
		}
		return *ppv ? S_OK : E_NOINTERFACE;
	}
	STDMETHOD_(ULONG, AddRef) () { return 1; }
	STDMETHOD_(ULONG, Release) () { return 1; }
	STDMETHOD(Indicate)(long numObjs, IWbemClassObject** ppObjs)
	{
		if(numObjs == 0 || ppObjs == NULL) return WBEM_E_INVALID_PARAMETER;
		for(long i = 0; i < numObjs; ++i)
		{
			if(IWbemClassObject* pProc = ppObjs[i])
			{
				HRESULT hr;
				VARIANT vName = {VT_EMPTY}, vCmdLine = {VT_EMPTY}, vPid = {VT_EMPTY};
				CIMTYPE type;
				hr = pProc->Get(L"Name", 0, &vName, &type, NULL);
				hr = pProc->Get(L"CommandLine", 0, &vCmdLine, &type, NULL);
				hr = pProc->Get(L"ProcessId", 0, &vPid, &type, NULL);
				printf("Found existing pid %lu, %ws (%ws)\n", V_UI4(&vPid), V_BSTR(&vName), V_BSTR(&vCmdLine));
				VariantClear(&vName);
				VariantClear(&vCmdLine);
				VariantClear(&vPid);
			}
		}
		return WBEM_S_NO_ERROR;
	}

	STDMETHOD(SetStatus)(long, HRESULT hr, BSTR param, IWbemClassObject*)
	{
		if(hr == WBEM_E_CALL_CANCELLED)
		{
			puts("Ending notifications for process enum");
		}
		else if(hr != WBEM_S_NO_ERROR)
		{
			printf(__FUNCTION__ " called with param=%ws, hr=%#08x\n", param, hr);
		}
		return WBEM_S_NO_ERROR;
	}
};

class ProcessChangeNotifier : public ProcessNotifier
{
	IWbemServices* pServices;
public:
	ProcessChangeNotifier(IWbemServices* pSvc) : pServices(pSvc) {}
	STDMETHOD(Indicate)(long numObjs, IWbemClassObject** ppObjs)
	{
		if(numObjs == 0 || ppObjs == NULL) return WBEM_E_INVALID_PARAMETER;
		// offset to --- is hardcoded below, change that if changing this
		// have to use Handle (aka the PID) to find the process, since it's
		// the 'key property', whatever that means
		STACK_BSTR(L"Win32_Process.Handle=\"------------", parentPath);
		for(long i = 0; i < numObjs; ++i)
		{
			if(IWbemClassObject* pProc = ppObjs[i])
			{
				HRESULT hr, hrExitCode;
				VARIANT vName, vPid, vExitCode;
				CIMTYPE type;
				hrExitCode = pProc->Get(L"ExitStatus", 0, &vExitCode, &type, NULL);
				hr = pProc->Get(L"ProcessName", 0, &vName, &type, NULL);
				hr = pProc->Get(L"ProcessId", 0, &vPid, &type, NULL);
				if(SUCCEEDED(hrExitCode))
				{
					// this is a stopped process
					printf("Process %ws (PID %lu) exited with code %lu\n", V_BSTR(&vName), V_UI4(&vPid), V_UI4(&vExitCode));
					VariantClear(&vExitCode);
				}
				else
				{
					// this is a created process
					VARIANT vParentId, vParentName = {VT_EMPTY};
					hr = pProc->Get(L"ParentProcessID", 0, &vParentId, &type, NULL);

					_snwprintf(parentPath.text + 22, 11, L"%lu\"", V_UI4(&vParentId));
					IWbemClassObject* pParentProcess = NULL;
					if(SUCCEEDED(hr = pServices->GetObject(parentPath, WBEM_FLAG_RETURN_WBEM_COMPLETE, NULL, &pParentProcess, NULL)))
					{
						hr = pParentProcess->Get(L"Name", 0, &vParentName, &type, NULL);
						pParentProcess->Release();
					}
					PCWSTR pParentName = (V_VT(&vParentName) == VT_BSTR) ? V_BSTR(&vParentName) : L"Unk";
					printf("Process %ws (%lu) created by process %ws\n", V_BSTR(&vName), V_UI4(&vPid), pParentName);
					VariantClear(&vParentName);
					VariantClear(&vParentId);
				}
				VariantClear(&vName);
				VariantClear(&vPid);
			}
		}
		return WBEM_S_NO_ERROR;
	}

	STDMETHOD(SetStatus)(long, HRESULT hr, BSTR param, IWbemClassObject*)
	{
		if(hr == WBEM_E_CALL_CANCELLED)
		{
			puts("Ending process start/stop notifications");
		}
		else if(hr != WBEM_S_NO_ERROR)
		{
			printf(__FUNCTION__ " called with param=%ws, hr=%#08x\n", param, hr);
		}
		return WBEM_S_NO_ERROR;
	}
};

int __cdecl main(int, char**)
{
	HRESULT hr;
	CHECK_HR_AND_RETURN(CoInitializeEx(NULL, COINIT_MULTITHREADED | COINIT_DISABLE_OLE1DDE));
	CHECK_HR_AND_RETURN(
		CoInitializeSecurity(
			NULL, 
			-1, 
			NULL,
			NULL,
			RPC_C_AUTHN_LEVEL_DEFAULT,
			RPC_C_IMP_LEVEL_IMPERSONATE,
			NULL,
			EOAC_DISABLE_AAA,
			NULL
		)
	);
	IWbemLocator* pLocator = NULL;
	CHECK_HR_AND_RETURN(CoCreateInstance(CLSID_WbemLocator, NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&pLocator))); 
	STACK_BSTR(L"root\\cimv2", cimName);
	IWbemServices* pNamespace = NULL;
	CHECK_HR_AND_RETURN(pLocator->ConnectServer(cimName, NULL, NULL, NULL, 0, NULL, NULL, &pNamespace));
	pLocator->Release();
	CHECK_HR_AND_RETURN(CoSetProxyBlanket(pNamespace, RPC_C_AUTHN_WINNT, RPC_C_AUTHN_NONE, NULL, RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE));
	puts("Press enter key to exit");

	STACK_BSTR(L"WQL", wql);
	STACK_BSTR(L"SELECT * From Win32_Process", processQuery);
	STACK_BSTR(L"SELECT * From Win32_ProcessStartTrace", newProcessQuery);
	STACK_BSTR(L"SELECT * From Win32_ProcessStopTrace", deadProcessQuery);
	ProcessNotifier procNotifier;
	CHECK_HR_AND_RETURN(pNamespace->ExecQueryAsync(wql, processQuery, 0, NULL, &procNotifier));
	ProcessChangeNotifier newProcNotifier(pNamespace);
	// these two will fail with WBEM_E_ACCESS_DENIED unless running as admin
	CHECK_HR(pNamespace->ExecNotificationQueryAsync(wql, newProcessQuery, 0, NULL, &newProcNotifier));
	CHECK_HR(pNamespace->ExecNotificationQueryAsync(wql, deadProcessQuery, 0, NULL, &newProcNotifier));
	getchar();
	pNamespace->CancelAsyncCall(&newProcNotifier);
	pNamespace->CancelAsyncCall(&procNotifier);
	pNamespace->Release();
	CoUninitialize();
}