//
// $Id: WMIProbe.cpp,v 1.2 2005/10/04 17:41:06 abuttner Exp $
//
//****************************************************************************************//
// Copyright (c) 2005, The MITRE Corporation
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification, are
// permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright notice, this list
//       of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright notice, this 
//       list of conditions and the following disclaimer in the documentation and/or other
//       materials provided with the distribution.
//     * Neither the name of The MITRE Corporation nor the names of its contributors may be
//       used to endorse or promote products derived from this software without specific 
//       prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY 
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
// OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT 
// SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
// OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
// TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
// EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
//****************************************************************************************//

#include "WMIProbe.h"

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Class WMIProbe  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

WMIProbe *WMIProbe::instance = NULL;

WMIProbe::WMIProbe()
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  Do nothing for now
	//
	// -----------------------------------------------------------------------
}

WMIProbe::~WMIProbe()
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  Do nothing for now
	//
	// -----------------------------------------------------------------------
}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Public Members  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

Probe* WMIProbe::Instance()
{

	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Ensure that the WMIProbe is a singleton.
	//
	// -----------------------------------------------------------------------

	// Use lazy initialization
	if(instance == NULL) 
		instance = new WMIProbe();

	return instance;	
}

pdVector WMIProbe::Run(ProbeData* probeDataIn)
{
	//------------------------------------------------------------------------------------//
	//
	//  ABSTRACT
	//
	//  Run the active directory probe. Return a vector of WMIData objects.
	//
	//------------------------------------------------------------------------------------//

	// Cast the input probe data object as an active directory probe data object

	WMIData* wmiProbeData = (WMIData*)probeDataIn;
	WMIData* tmpProbeData;
	pdVector instanceVector;
	pdVector returnVector;

	// If the naming_context, relative_dn, and attribute elements are of a litteral type,
	// we do not need to do any of the pattern matching stuff.

	try
	{
		if((wmiProbeData->wmi_namespace)->type == literal &&
		   (wmiProbeData->wql)->type == literal)
		{
			instanceVector = GetWMI(wmiProbeData->wmi_namespace->object,
								    wmiProbeData->wql->object);

			// Iterate through instanceVector and for each item, add it to the returnVector.

			pdVector::iterator instances;
			for (instances=instanceVector.begin(); instances!=instanceVector.end(); instances++)
			{
				returnVector.push_back((*instances));
			}
		}
		else
		{
			// At least one of the elements have a pattern matching type.  Each pattern match must
			// be evaluated and for each possible wmi object that satisfies the pattern
			// match, the associated data must be retreived.

			string errorMessage = "";

			if(wmiProbeData->wmi_namespace->type == pattern_match)
			{
				errorMessage.append("(WMIProbe) The Definition Interpreter does not support a regular expression in the Namespace field.");
			}
			else if(wmiProbeData->wql->type == pattern_match)
				errorMessage.append("(WMIProbe) The Definition Interpreter does not support a regular expression in the WQL field.");

			throw WMIProbeException(errorMessage, ERROR_FATAL);
		}

	}
	catch(Exception ex)
	{
		tmpProbeData = new WMIData();
		tmpProbeData->SetMessage(ex.GetErrorMessage());
		tmpProbeData->SetStatus(error);
		tmpProbeData->wmi_namespace->object = wmiProbeData->wmi_namespace->object;
		tmpProbeData->wmi_namespace->type = wmiProbeData->wmi_namespace->type;
		tmpProbeData->wql->object = wmiProbeData->wql->object;
		tmpProbeData->wql->type = wmiProbeData->wql->type;
		tmpProbeData->isPatternMatchObject = wmiProbeData->isPatternMatchObject;
		returnVector.push_back(tmpProbeData);
		tmpProbeData = NULL;

	}
	catch(...)
	{
		tmpProbeData = new WMIData();
		tmpProbeData->SetMessage("Error: (WMIProbe) Unknown exception occured while running the wmi probe");
		tmpProbeData->SetStatus(error);
		tmpProbeData->wmi_namespace->object = wmiProbeData->wmi_namespace->object;
		tmpProbeData->wmi_namespace->type = wmiProbeData->wmi_namespace->type;
		tmpProbeData->wql->object = wmiProbeData->wql->object;
		tmpProbeData->wql->type = wmiProbeData->wql->type;
		tmpProbeData->isPatternMatchObject = wmiProbeData->isPatternMatchObject;
		returnVector.push_back(tmpProbeData);
		tmpProbeData = NULL;
	}

	return returnVector;
}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Private Members  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

pdVector WMIProbe::GetWMI(string wmi_namespaceIn, string wqlIn)
{
	//------------------------------------------------------------------------------------//
	//
	//  ABSTRACT
	//
	//  Gather new data and put in a WMIData object.
	//
	//------------------------------------------------------------------------------------//

	HRESULT hres;

	// Initialize the output structure, an WMIData object.

	WMIData* resultData;
	pdVector resultVector;

	IWbemLocator *pLoc = NULL;
	IWbemServices *pSvc = NULL;
	IEnumWbemClassObject* pEnumerator = NULL;
	IWbemClassObject *pclsObj[1];

	ULONG uReturn = 0;

	try {
		// establish COM connection
		hres = CoInitializeEx(NULL, COINIT_APARTMENTTHREADED); 
		if (FAILED(hres))
		{
			string errorMessage = _com_error(hres).ErrorMessage();
			throw WMIProbeException("(WMIProbe) Failed to initialize COM library.  " + errorMessage, ERROR_FATAL);
		}

		// set security of COM connection to the default
		hres =  CoInitializeSecurity(NULL, -1, NULL, NULL, RPC_C_AUTHN_LEVEL_DEFAULT, RPC_C_IMP_LEVEL_IMPERSONATE, NULL, EOAC_NONE, NULL);
		if (FAILED(hres))
		{

			string errorMessage = _com_error(hres).ErrorMessage();
			throw WMIProbeException("(WMIProbe) Failed to initialize COM security.  " + errorMessage, ERROR_FATAL);
		}

		// find the WMI Locator
		hres = CoCreateInstance(CLSID_WbemLocator, NULL, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID *) &pLoc);
	    if (FAILED(hres))
		{
			string errorMessage = _com_error(hres).ErrorMessage();
			throw WMIProbeException("(WMIProbe) Failed to create IWbemLocator object.  " + errorMessage, ERROR_FATAL);
	    }

		// Connect to the specified namespace with the current user.
		hres = pLoc->ConnectServer(_bstr_t(wmi_namespaceIn.c_str()), NULL, NULL, 0, NULL, 0, 0, &pSvc);
	    if (FAILED(hres))
		{
			string errorMessage = _com_error(hres).ErrorMessage();
			throw WMIProbeException("(WMIProbe) Unable to connect to the '" + wmi_namespaceIn + "' namespace.  " + errorMessage, ERROR_FATAL);
	    }
		// connected to WMI

		// setting the security levels of the WMI connection
		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 (FAILED(hres))
		{
			string errorMessage = _com_error(hres).ErrorMessage();   
 			throw WMIProbeException("(WMIProbe) Unable to set the WMI proxy blanket.  " + errorMessage, ERROR_FATAL);
		}

		// run the query
		hres = pSvc->ExecQuery(_bstr_t(L"WQL"), _bstr_t(wqlIn.c_str()), WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY, NULL, &pEnumerator);
        if (FAILED(hres))
		{
 			throw WMIProbeException("(WMIProbe) Wmi query failed. ('" + wqlIn + "')", ERROR_FATAL);
	    }


		// retrieve the result
		HRESULT    enumhRes = WBEM_S_NO_ERROR;
	    while (enumhRes == WBEM_S_NO_ERROR)
		{
			resultData = new WMIData();
			resultData->wmi_namespace->object = wmi_namespaceIn;
			resultData->wql->object = wqlIn;

			enumhRes = pEnumerator->Next(WBEM_INFINITE, 1, pclsObj, &uReturn); // iterate through each instance returned
			if(uReturn == 0)
			{
				// record the result
				resultData->SetMessage("(WMIProbe) The wmi query did not produce any results.");
				resultData->SetStatus(doesNotExist);
			}
			else
			{
				VARIANT vtProp;
				VariantInit(&vtProp);

				// Get the value of the Name property
				
				string fieldName = "unknown error";
				if (GetWqlField(wqlIn, &fieldName) == true)
				{
					hres = pclsObj[0]->Get(bstr_t(fieldName.c_str()), 0, &vtProp, 0, 0);

					// Convert the UNICODE name to a char name.

					char* szChar = NULL;
					size_t size = 0;
					string strFieldName = "";

					if((size = wcstombs(0, vtProp.bstrVal, 0)) != -1)
					{
						szChar = new char[size + 1];
						szChar[size] = NULL;
						wcstombs(szChar, vtProp.bstrVal, size);
					}

					if (size == -1)
					{
						VariantClear(&vtProp);
						throw WMIProbeException("(WMIProbe) There was an error converting the UNICODE field name.", ERROR_FATAL);
					}

					strFieldName = szChar;

					// record the result
					resultData->result->value = strFieldName;
					resultData->result->dataType = stringType;
				}
				else
				{
					VariantClear(&vtProp);
					throw WMIProbeException("(WMIProbe) Invalid field in WQL statement.  ERROR MESSAGE - " + fieldName, ERROR_FATAL);
				}

				VariantClear(&vtProp);
			}

			resultVector.push_back(resultData);

			// for version 4.1 we only get a single result so we can break now.
			break;
		}

	}
	catch(WMIProbeException ex)
	{
		resultData = new WMIData();
		resultData->wmi_namespace->object = wmi_namespaceIn;
		resultData->wql->object = wqlIn;
		resultData->SetMessage(resultData->GetMessage() + " Message: " + ex.GetErrorMessage());
		resultData->SetStatus(error);
		resultVector.push_back(resultData);
	}
	catch(...)
	{	
		string errMsg = "";
		errMsg.append("(WMIProbe) Unknown error attempting to get wmi ");
		errMsg.append("information for the namespace '");
		errMsg.append(wmi_namespaceIn);
		errMsg.append("' and the wql query '");
		errMsg.append(wqlIn);
		errMsg.append("'\n");

		resultData = new WMIData();
		resultData->wmi_namespace->object = wmi_namespaceIn;
		resultData->wql->object = wqlIn;
		resultData->SetMessage(resultData->GetMessage() + " Error Message: " + errMsg);
		resultData->SetStatus(error);
		resultVector.push_back(resultData);
	}

	// clean-up
	for ( ULONG n = 0; n < uReturn; n++ ) pclsObj[n]->Release();
	if (pEnumerator != NULL) pEnumerator->Release();
	if (pSvc != NULL) pSvc->Release();
	if (pLoc != NULL) pLoc->Release();

	// Close the COM library on the current thread
	CoUninitialize();

	return resultVector;
}

bool WMIProbe::GetWqlField(string wqlIn, string* fieldName)
{
	//------------------------------------------------------------------------------------//
	//
	//  ABSTRACT
	//
	//  Parse the WQL and extract the field in the select statement.  If we get the field,
	//  return true and fieldName is set to the name of the field.  If we cannot get the
	//  field, return false and fieldName is set to an error message
	//
	//------------------------------------------------------------------------------------//

	int endSelect, startFrom;

	// Create a copy of the WQL string in upper case.  This way we can find the opening
	// SELECT and FROM words which are not case sensitive.

	string wqlUpperCase = wqlIn;
	for(unsigned int i=0;i<wqlIn.length();i++)
	{
		wqlUpperCase[i] = toupper(wqlIn[i]);
	}

	// Find the opening SELECT statement

	endSelect = wqlUpperCase.find("SELECT ",0);
	if (endSelect == -1)
	{
		*fieldName = "While searching for the fieldname, couldn't find the opening SELECT.";
		return false;
	}
	endSelect = endSelect + 7;

	// Find the FROM statement.

	startFrom = wqlUpperCase.find(" FROM", endSelect);
	if (endSelect == -1)
	{
		*fieldName = "While searching for the fieldname, couldn't find the FROM statement.";
		return false;
	}

	*fieldName = wqlIn.substr(endSelect, startFrom-endSelect);
	return true;
}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~  Class WMIProbeException  ~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//

WMIProbeException::WMIProbeException(string errMsgIn, int severity) : Exception(errMsgIn, severity)
{
}

WMIProbeException::~WMIProbeException()
{
}
