// **************************************************************************
//
//	CAdvClientDlg::OnDiskdetails()
//
// Description:
//		Enumerates the properties of the C: drive using the 'GetNames()'
//		technique. The technique uses safearrays.
//
// Parameters:
//		None.
//
// Returns:
//		Nothing.
//
// Globals accessed:
//		None.
//
// Globals modified:
//		None.
//
//===========================================================================
void CAdvClientDlg::OnDiskdetails() 
{
	HRESULT  hRes;
	long lLower, lUpper, lCount; 
	SAFEARRAY *psaNames = NULL;
	BSTR PropName = NULL;
	VARIANT varString, pVal;
	WCHAR *pBuf;
	CString clMyBuff;

	IWbemClassObject *pDriveInst = NULL;
	IWbemQualifierSet *pQualSet = NULL;

	VariantInit(&varString);
	VariantInit(&pVal);

	m_outputList.ResetContent();
	m_outputList.AddString(_T("working..."));

	//-------------------------------
	// Get the instance for C: drive.
	BSTR driveName = SysAllocString(L"Win32_LogicalDisk.DeviceID=\"C:\"");
	if (!driveName)
	{
		TRACE(_T("SysAllocString failed: not enough memory\n"));
		return;
	}
	BSTR cimType = SysAllocString(L"CIMTYPE");
	if (!cimType)
	{
		SysFreeString(driveName);
		TRACE(_T("SysAllocString failed: not enough memory\n"));
		return;
	}		
	BSTR keyQual = SysAllocString(L"key");
	if (!keyQual)
	{
		SysFreeString(driveName);
		SysFreeString(cimType);
		TRACE(_T("SysAllocString failed: not enough memory\n"));
		return;
	}

    if((hRes = m_pIWbemServices->GetObject(driveName,
										0L,
										NULL,
										&pDriveInst,
										NULL)) == S_OK)
	{

		m_outputList.ResetContent();

		//-------------------------------
		// Get the property names
		if((hRes = pDriveInst->GetNames(NULL, 
										WBEM_FLAG_ALWAYS | 
										WBEM_FLAG_NONSYSTEM_ONLY, 
										NULL, 
										&psaNames)) == S_OK)
		{
			//-------------------------------
			// Get the upper and lower bounds of the Names array
			if((hRes = SafeArrayGetLBound(psaNames, 1, &lLower)) != S_OK) 
			{
				TRACE(_T("Couldn't get safe array lbound\n"));
				SafeArrayDestroy(psaNames);
				return;
			}

			//-------------------------------
			if((hRes = SafeArrayGetUBound(psaNames, 1, &lUpper)) != S_OK) 
			{
				TRACE(_T("Couldn't get safe array ubound\n"));
				SafeArrayDestroy(psaNames);
				return;
			}

		
			//-------------------------------
			// For all properties...
			for (lCount = lLower; lCount <= lUpper; lCount++) 
			{
				//-----------------------------------------------
				// I'm formatting each property as:
				//   name (type) ==> value
				//-----------------------------------------------

				//-------------------------------
				// get the property name for this element
				if((hRes = SafeArrayGetElement(psaNames, 
												&lCount, 
												&PropName)) == S_OK)
				{
					clMyBuff = PropName;

					// print variable type for property value
					clMyBuff += _T(" (");

					// Get pointer to property qualifiers
					// this mess is due to the fact that system properties don't have qualifiers
					if ((pDriveInst->GetPropertyQualifierSet(PropName, &pQualSet)) == S_OK) 
					{
						// Get and print syntax attribute (if any)
						if ((pQualSet->Get(cimType, 0L, &pVal, NULL)) == S_OK) 
						{
						   clMyBuff += V_BSTR(&pVal);
						} 
						else if (hRes != WBEM_E_NOT_FOUND) 
						{  // some other error
						   TRACE(_T("Could not get syntax qualifier\n"));
						   break;
						}
						VariantClear(&pVal);

						//-------------------------------
						// If this is a key field, print an asterisk
						if(((hRes = pQualSet->Get(keyQual, 
												0L, 
												&pVal, 
												NULL)) == S_OK) && 
							(pVal.boolVal))
						{ // Yes, it's a key
						   clMyBuff += _T(")*");
						} 
						else if (hRes == WBEM_E_NOT_FOUND) 
						{  // not a key qualifier
						   clMyBuff += _T(")");
						} 
						else 
						{ // some other error
						   TRACE(_T("Could not get key qualifier\n"));
						   break;
						}
						// done with the qualifierSet.
						if (pQualSet)
						{ 
							pQualSet->Release(); 
							pQualSet = NULL;
						}
					} 
					else 
					{
						clMyBuff += _T(")");
					} //endif pDriveClass->GetPropertyQualifierSet()

					//-------------------------------
					// Get the value for the property.
					if((hRes = pDriveInst->Get(PropName, 
												0L, 
												&varString, 
												NULL, NULL)) == S_OK) 
					{
						// Print the value
						clMyBuff += _T("   ==> ");
						clMyBuff += ValueToString(&varString, &pBuf);
						
						m_outputList.AddString(clMyBuff);

						free(pBuf); // allocated by ValueToString()
					}
					else
					{
						TRACE(_T("Couldn't get Property Value\n"));
						break;
					} //endif pDriveClass->Get()

					VariantClear(&varString);
					VariantClear(&pVal);
				}
				else // SafeArrayGetElement() failed
				{
					TRACE(_T("Couldn't get safe array element\n"));
					break;
				} //endif SafeArrayGetElement()

			} // endfor

			// cleanup.
			SysFreeString(PropName);
			SysFreeString(keyQual);
			SysFreeString(cimType);
			SafeArrayDestroy(psaNames);
			VariantClear(&varString);
			VariantClear(&pVal);
		}
		else // pDriveClass->GetNames() failed
		{
			TRACE(_T("Couldn't GetNames\n"));
		} //endif pDriveClass->GetNames()

		// done with drive instance.
		if (pDriveInst)
		{ 
			pDriveInst->Release(); 
			pDriveInst = NULL;
		}
	} //endif GetObject()
}
/*!
 * @brief Perform a WMI query.
 * @param lpwRoot Name of the root object that is to be queried against.
 * @param lpwQuery The filter to use when reading objects (LDAP style).
 * @param response The response \c Packet to add the results to.
 */
DWORD wmi_query(LPCWSTR lpwRoot, LPWSTR lpwQuery, Packet* response)
{
	HRESULT hResult;

	dprintf("[WMI] Initialising COM");
	if ((hResult = CoInitializeEx(NULL, COINIT_MULTITHREADED)) == S_OK)
	{
		dprintf("[WMI] COM initialised");
		IWbemLocator* pLocator = NULL;
		IWbemServices* pServices = NULL;
		IEnumWbemClassObject* pEnumerator = NULL;
		IWbemClassObject* pSuperClass = NULL;
		IWbemClassObject* pObj = NULL;
		Tlv* valueTlvs = NULL;
		char* values = NULL;
		VARIANT** fields = NULL;

		do
		{
			if (FAILED(hResult = CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, 0)))
			{
				dprintf("[WMI] Failed to initialize security: %x", hResult);
				break;
			}
			dprintf("[WMI] Security initialised");

			if (FAILED(hResult = CoCreateInstance(CLSID_WbemLocator, NULL, CLSCTX_ALL, IID_PPV_ARGS(&pLocator))))
			{
				dprintf("[WMI] Failed to create WbemLocator: %x", hResult);
				break;
			}
			dprintf("[WMI] WbemLocator created.");

			if (FAILED(hResult = pLocator->ConnectServer(_bstr_t(lpwRoot), NULL, NULL, NULL, WBEM_FLAG_CONNECT_USE_MAX_WAIT, NULL, NULL, &pServices)))
			{
				dprintf("[WMI] Failed to create WbemServices at %S: %x", lpwRoot, hResult);
				break;
			}
			dprintf("[WMI] WbemServices created.");

			if (FAILED(hResult = pServices->ExecQuery(L"WQL", lpwQuery, WBEM_FLAG_FORWARD_ONLY, NULL, &pEnumerator)))
			{
				dprintf("[WMI] Failed to create Enumerator for query %S: %x", lpwQuery, hResult);
				break;
			}
			dprintf("[WMI] Enumerated created.");

			ULONG numFound;
			if (FAILED(hResult = pEnumerator->Next(ENUM_TIMEOUT, 1, &pObj, &numFound)))
			{
				dprintf("[WMI] Failed to get the first query element: %x", lpwQuery, hResult);
				break;
			}
			dprintf("[WMI] First result read. hr=%x p=%p", hResult, pObj);

			if (hResult == WBEM_S_FALSE)
			{
				// this is not an error
				dprintf("[WMI] No results found!");
				break;
			}

			// get the names of the fields out of the first object before doing anything else.
			LPSAFEARRAY pFieldArray = NULL;
			if (FAILED(hResult = pObj->GetNames(NULL, WBEM_FLAG_ALWAYS, NULL, &pFieldArray)))
			{
				dprintf("[WMI] Failed to get field names: %x", hResult);
				break;
			}
			dprintf("[WMI] Field Names extracted. hr=%x p=%p", hResult, pFieldArray);

			// lock the array
			if (FAILED(hResult = SafeArrayLock(pFieldArray)))
			{
				dprintf("[WMI] Failed to get array dimension: %x", hResult);
				break;
			}
			dprintf("[WMI] Field name array locked.");

			do
			{
				dprintf("[WMI] Array dimensions: %u", SafeArrayGetDim(pFieldArray));

				// this array is just one dimension, let's get the bounds of the first dimension
				LONG lBound, uBound;
				if (FAILED(hResult = SafeArrayGetLBound(pFieldArray, 1, &lBound))
					|| FAILED(hResult = SafeArrayGetUBound(pFieldArray, 1, &uBound)))
				{
					dprintf("[WMI] Failed to get array dimensions: %x", hResult);
					break;
				}
				dprintf("[WMI] Bounds: %u to %u", lBound, uBound);

				LONG fieldCount = uBound - lBound - SYSTEM_FIELD_COUNT - 1;
				dprintf("[WMI] Query results in %u fields", fieldCount);

				fields = (VARIANT**)malloc(fieldCount * sizeof(VARIANT**));
				valueTlvs = (Tlv*)malloc(fieldCount * sizeof(Tlv));
				values = (char*)malloc(fieldCount * FIELD_SIZE);
				memset(fields, 0, fieldCount * sizeof(VARIANT**));
				memset(valueTlvs, 0, fieldCount * sizeof(Tlv));
				memset(values, 0, fieldCount * FIELD_SIZE);

				for (LONG i = 0; i < fieldCount; ++i)
				{
					LONG indices[1] = { i + SYSTEM_FIELD_COUNT };
					char* fieldName = values + (i * FIELD_SIZE);
					SafeArrayPtrOfIndex(pFieldArray, indices, (void**)&fields[i]);
					_bstr_t bstr(fields[i]->bstrVal);

					strncpy_s(fieldName, FIELD_SIZE, (const char*)bstr, FIELD_SIZE - 1);

					valueTlvs[i].header.type = TLV_TYPE_EXT_WMI_FIELD;
					valueTlvs[i].header.length = (UINT)strlen(fieldName) + 1;
					valueTlvs[i].buffer = (PUCHAR)fieldName;

					dprintf("[WMI] Added header field: %s", fieldName);
				}

				dprintf("[WMI] added all field headers");
				// add the field names to the packet
				packet_add_tlv_group(response, TLV_TYPE_EXT_WMI_FIELDS, valueTlvs, fieldCount);

				dprintf("[WMI] processing values...");
				// with that horrible pain out of the way, let's actually grab the data
				do
				{
					if (FAILED(hResult))
					{
						dprintf("[WMI] Loop exited via %x", hResult);
						break;
					}

					memset(valueTlvs, 0, fieldCount * sizeof(Tlv));
					memset(values, 0, fieldCount * FIELD_SIZE);

					for (LONG i = 0; i < fieldCount; ++i)
					{
						char* value = values + (i * FIELD_SIZE);
						valueTlvs[i].header.type = TLV_TYPE_EXT_WMI_VALUE;
						valueTlvs[i].buffer = (PUCHAR)value;

						VARIANT varValue;
						VariantInit(&varValue);

						_bstr_t field(fields[i]->bstrVal);
						dprintf("[WMI] Extracting value for %s", (char*)field);
						if (SUCCEEDED(pObj->Get(field, 0, &varValue, NULL, NULL)))
						{
							variant_to_string(_variant_t(varValue), value, FIELD_SIZE);
						}

						valueTlvs[i].header.length = (UINT)strlen(value) + 1;
						dprintf("[WMI] Added value for %s: %s", (char*)_bstr_t(fields[i]->bstrVal), value);
					}

					// add the field values to the packet
					packet_add_tlv_group(response, TLV_TYPE_EXT_WMI_VALUES, valueTlvs, fieldCount);

					pObj->Release();
					pObj = NULL;
				} while ((hResult = pEnumerator->Next(ENUM_TIMEOUT, 1, &pObj, &numFound)) != WBEM_S_FALSE);

			} while (0);

			SafeArrayUnlock(pFieldArray);
		} while (0);

		if (fields)
		{
			free(fields);
		}

		if (values)
		{
			free(values);
		}

		if (valueTlvs)
		{
			free(valueTlvs);
		}

		if (pObj)
		{
			pObj->Release();
		}

		if (pEnumerator)
		{
			pEnumerator->Release();
		}

		if (pServices)
		{
			pServices->Release();
		}

		if (pLocator)
		{
			pLocator->Release();
		}
		CoUninitialize();

		if (SUCCEEDED(hResult))
		{
			hResult = S_OK;
			dprintf("[WMI] Things appeard to go well!");
		}
	}
	else
	{
		dprintf("[WMI] Failed to initialize COM");
	}

	if (FAILED(hResult))
	{
		// if we failed, we're going to convert the error to a string, add it and still return success, but we'll
		// also include the hresult.
		char errorMessage[1024];
		memset(errorMessage, 0, 1024);
		_com_error comError(hResult);
		_snprintf_s(errorMessage, 1024, 1023, "%s (0x%x)", comError.ErrorMessage(), hResult);
		dprintf("[WMI] returning error -> %s", errorMessage);
		packet_add_tlv_string(response, TLV_TYPE_EXT_WMI_ERROR, errorMessage);
		hResult = S_OK;
	}

	return (DWORD)hResult;
}