/*
** Zabbix
** Copyright (C) 2001-2014 Zabbix SIA
**
** This program is free software; you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation; either version 2 of the License, or
** (at your option) any later version.
**
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with this program; if not, write to the Free Software
** Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
**/

#include "sysinc.h"

extern "C"
{
#	include "common.h"
#	include "sysinfo.h"
#	include "log.h"
}

#include <comdef.h>
#include <Wbemidl.h>

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


struct CoInitializeGuard
{
	bool initialized;

	CoInitializeGuard() : initialized(false)
	{
		HRESULT hres =  CoInitializeEx(0, COINIT_MULTITHREADED);
		if (SUCCEEDED(hres)) initialized = true;
	}

	~CoInitializeGuard()
	{
		if (initialized) {
			CoUninitialize();
			initialized = false;
		}
	}
};

struct CoInitializeSecurityGuard
{
	bool initialized;

	CoInitializeSecurityGuard() : initialized(false)
	{
		HRESULT hres =  CoInitializeSecurity(
			NULL,    
			-1,               
			NULL,
			NULL,
			RPC_C_AUTHN_LEVEL_DEFAULT,
			RPC_C_IMP_LEVEL_IMPERSONATE,
			NULL,
			EOAC_NONE,
			NULL
		);
  
		if (SUCCEEDED(hres)) initialized = true;
     }
  
	~CoInitializeSecurityGuard()
     {
		 initialized = false;
     }
};

struct CoCreateInstanceGuard
{
	IWbemLocator *pLoc;
  
	CoCreateInstanceGuard() : pLoc(0)
	{
		HRESULT hres = CoCreateInstance(
			CLSID_WbemLocator,            
			0,
			CLSCTX_INPROC_SERVER,
			IID_IWbemLocator, 
			(LPVOID *) &pLoc);
	}
  
	~CoCreateInstanceGuard()
	{
		if (pLoc != 0) {
			pLoc->Release();
		}
	}
};

struct ConnectServerGuard
{
	IWbemServices *pSvc;
	bool initialized;

	ConnectServerGuard(IWbemLocator *pLoc, char *strWMINameSpace) : pSvc(0), initialized(false)
	{
		HRESULT hres = pLoc->ConnectServer(
			_bstr_t(strWMINameSpace),
			NULL,                    
			NULL,                    
			0,                       
			NULL,                    
			0,                       
			0,                       
			&pSvc                    
		);                             
	
	}

	~ConnectServerGuard()
	{
		if (pSvc != 0) {
			pSvc->Release();
		}
	}
};

struct CoSetProxyBlanketGuard
{
	bool initialized;

	CoSetProxyBlanketGuard(IWbemServices *pSvc)
	{
		HRESULT hres = CoSetProxyBlanket(
			pSvc,                         
			RPC_C_AUTHN_WINNT,            
			RPC_C_AUTHZ_NONE,             
			NULL,                         
			RPC_C_AUTHN_LEVEL_CALL,       
			RPC_C_IMP_LEVEL_IMPERSONATE,  
			NULL,                         
			EOAC_NONE                     
		);

    if (SUCCEEDED(hres)) initialized = true;

	}

	~CoSetProxyBlanketGuard()
	{
		initialized = false;
	}
};

static	CoInitializeGuard cig;
static	CoInitializeSecurityGuard cisg;

HRESULT ExecQuery(char* snamespace, char* squery, VARIANT* pvout)
{
	HRESULT		hres;
	ULONG		uReturn = 0;
	ULONG		uTimeout = 60000; //60 sec timeout

	if (cig.initialized && cisg.initialized)
	{
		zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: connected to COM");
		CoCreateInstanceGuard ccig;
		
		if ( ccig.pLoc != 0 )
		{
			ConnectServerGuard csg(ccig.pLoc,snamespace);
			CoSetProxyBlanketGuard cspb(csg.pSvc);
	
			if (csg.pSvc != 0 && cspb.initialized)
			{
				zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: connected to namepace");
	
				IEnumWbemClassObject*		pEnumerator = 0;
	
				hres = csg.pSvc->ExecQuery(_bstr_t("WQL"), _bstr_t(squery),
					WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumerator);
	
				if ( hres == WBEM_S_NO_ERROR ) 
				{
					zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: query successfull");
	
					IWbemClassObject*			pclsObj = 0;
	
					hres = pEnumerator->Next(uTimeout, 1, &pclsObj, &uReturn);
					
					if (uReturn != 0 && hres == WBEM_S_NO_ERROR)
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: successfull get WMI object");
						hres = pclsObj->BeginEnumeration(WBEM_FLAG_NONSYSTEM_ONLY);
						if ( hres == WBEM_S_NO_ERROR )
						{
							zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: successfull start WMI object properties enumeration");
							hres = pclsObj->Next(0, NULL, pvout, 0, 0);
							if ( hres != WBEM_S_NO_ERROR )
							{
								zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: fail to get WMI object property");
								pclsObj->Release();
								pEnumerator->Release();
								return hres;
							} else 
							{
								zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: sucessfull retrieved WMI value");
	
								pclsObj->EndEnumeration();
								pclsObj->Release();
								pEnumerator->Release();
								return S_OK;
	
							}
						}
						else
						{
							zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: fail to start WMI object properties enumeration");
							pclsObj->Release();
							pEnumerator->Release();
							return hres;
						}
					}
					else
					{	
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: fail to get WMI object, empty result or timeout");
						pEnumerator->Release();
						return hres;
					}
				}
				else
				{
					zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: failed to execute WMI query %s", squery);
					return hres;
				}
	
			} else
			{
				zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot connect to WMI namespace");
				return E_FAIL;
			}
		} else
		{
			zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot get WMI locator");
			return E_FAIL;
		}
	} else
	{
		zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: COM not initialized");
		return E_FAIL;
	}
	
}

extern "C" void	zbx_co_uninitialize(){}

extern "C" int	WMI_GET(AGENT_REQUEST *request, AGENT_RESULT *result)
{
	char			*wmi_namespace, *wmi_query;
	VARIANT			vtProp;
	HRESULT			hres;
	int			ret = SYSINFO_RET_FAIL;

	if (2 != request->nparam)
	{
		SET_MSG_RESULT(result, zbx_strdup(NULL, "Invalid number of parameters."));
		return SYSINFO_RET_FAIL;
	}

	wmi_namespace = get_rparam(request, 0);
	wmi_query = get_rparam(request, 1);

	wmi_query = string_replace( wmi_query, "\\", "\\\\"); 

	VariantInit(&vtProp);

	hres = ExecQuery(wmi_namespace,wmi_query, &vtProp);

	if (SUCCEEDED(hres))
	{

		if (0 != (vtProp.vt & VT_ARRAY))
		{
			SAFEARRAY *safearray = vtProp.parray;
			VARTYPE vtype;
			LONG lIndex = 0;

			SafeArrayGetVartype(safearray, &vtype);

			switch(vtype)
			{
				case VT_EMPTY:
				case VT_NULL:
					break;
				case VT_I8:
				case VT_I4:
				case VT_UI1:
				case VT_I2:
				case VT_I1:
				case VT_UI2:
				case VT_UI4:
				case VT_UI8:
				case VT_INT:
				case VT_UINT:

					unsigned int uint;

					SafeArrayGetElement(safearray, &lIndex, &uint);
					if (FAILED(hres))
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot get unsigned integer from array");
						break;
					}
					SET_UI64_RESULT(result, uint);
					ret = SYSINFO_RET_OK;
					break;

				case VT_R4:
				case VT_R8:

					double vt_r8;
					hres = SafeArrayGetElement(safearray, &lIndex, &vt_r8);
					
					if (FAILED(hres))
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot get double from array");
						break;
					}

					SET_DBL_RESULT(result, vt_r8);
					ret = SYSINFO_RET_OK;
					break;


				default:

					BSTR bstr;
					
					hres = SafeArrayGetElement(safearray, &lIndex, &bstr);
					if (FAILED(hres))
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot get BSTR from array");
						break;
					}

					SET_TEXT_RESULT(result, zbx_strdup(NULL, (char*)_bstr_t(bstr)));
					ret = SYSINFO_RET_OK;
					SysFreeString(bstr);
					break;

			}		
		} 
		else
		{
			switch(vtProp.vt)
			{
				case VT_EMPTY:
				case VT_NULL:
					break;
				case VT_I8:
				case VT_I4:
				case VT_UI1:
				case VT_I2:
				case VT_I1:
				case VT_UI2:
				case VT_UI4:
				case VT_UI8:
				case VT_INT:
				case VT_UINT:

					hres = VariantChangeType(&vtProp, &vtProp, 0, VT_I8);
					if (FAILED(hres))
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot convert WMI result of type %d to VT_I8", vtProp.vt);
						break;
					}
					SET_UI64_RESULT(result, vtProp.llVal);
					ret = SYSINFO_RET_OK;
					break;

				case VT_R4:
				case VT_R8:

					hres = VariantChangeType(&vtProp, &vtProp, 0, VT_R8);
					if (FAILED(hres))
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot convert WMI result of type %d to VT_R8", vtProp.vt);
						break;
					}

					SET_DBL_RESULT(result, vtProp.dblVal);
					ret = SYSINFO_RET_OK;
					break;

				default:

					hres = VariantChangeType(&vtProp, &vtProp, VARIANT_ALPHABOOL, VT_BSTR);
					if (FAILED(hres))
					{
						zabbix_log(LOG_LEVEL_DEBUG, "wmi.get: cannot convert WMI result of type %d to VT_BSTR", vtProp.vt);
						break;
					}

					SET_TEXT_RESULT(result, zbx_strdup(NULL, (char*)_bstr_t(vtProp.bstrVal)));
					ret = SYSINFO_RET_OK;
			}
		}
	}

	VariantClear(&vtProp); 

	if (SYSINFO_RET_FAIL == ret)
	{
		zabbix_log(LOG_LEVEL_ERR, "wmi.get: Query failed %s %s", wmi_namespace, wmi_query);
		SET_MSG_RESULT(result, zbx_strdup(NULL, "Cannot obtain WMI information.")); 
	}

	return ret;
}