2025-04-27 07:49:33 -04:00

430 lines
10 KiB
C++

// EssTest.cpp
#include "stdafx.h"
#include "EssTest.h"
// Our single global instance.
CEssTest g_essTest;
CEssTest::CEssTest() :
m_logLevel(LOGLEVEL_RESULTS_AND_ERRORS),
m_bKeepLogs(FALSE)
{
InitializeCriticalSection(&m_csOutput);
}
CEssTest::~CEssTest()
{
Cleanup();
}
DWORD WINAPI CEssTest::RunItemProc(CWorkItem *pItem)
{
pItem->Run();
return 0;
}
void CEssTest::Run()
{
if (SUCCEEDED(Init()))
{
int nItems = m_listItems.size(),
i = 0;
HANDLE *phthreadItems = new HANDLE[nItems];
for (CWorkItemListIterator item = m_listItems.begin();
item != m_listItems.end();
item++)
{
DWORD dwID;
phthreadItems[i++] =
CreateThread(
NULL,
0,
(LPTHREAD_START_ROUTINE) RunItemProc,
*item,
0,
&dwID);
}
WaitForMultipleObjects(nItems, phthreadItems, TRUE, INFINITE);
for (i = 0; i < nItems; i++)
CloseHandle(phthreadItems[i]);
delete phthreadItems;
}
}
void CEssTest::Pause()
{
}
void CEssTest::Stop()
{
}
HRESULT CEssTest::LoadWorkItems()
{
IEnumWbemClassObject *pEnum = NULL;
_bstr_t strClass = L"MSFT_EssTestWorkItem";
HRESULT hr;
hr =
m_pDefNamespace->CreateInstanceEnum(
strClass,
WBEM_FLAG_DEEP | WBEM_FLAG_RETURN_IMMEDIATELY |
WBEM_FLAG_FORWARD_ONLY,
NULL,
&pEnum);
if (SUCCEEDED(hr))
{
IWbemClassObject *pObj = NULL;
DWORD nCount,
dwID = 0;
while(SUCCEEDED(hr = pEnum->Next(WBEM_INFINITE, 1, &pObj, &nCount)) &&
nCount > 0)
{
_variant_t vEnabled(false);
_variant_t vName(L"???");
pObj->Get(L"Name", 0, &vName, NULL, NULL);
pObj->Get(L"Enabled", 0, &vEnabled, NULL, NULL);
if ((bool) vEnabled == true)
{
CWorkItem *pItem = new CWorkItem;
HRESULT hr;
if (SUCCEEDED(hr = pItem->Init(pObj, dwID)))
{
m_listItems.push_back(pItem);
PrintStatus("Loaded enabled work item '%S'.", V_BSTR(&vName));
dwID++;
}
else
{
PrintError(
"Failed to init work item '%S': 0x%X",
V_BSTR(&vName),
hr);
}
pObj->Release();
}
else
{
PrintStatus("Skipping disabled work item '%S'.", V_BSTR(&vName));
}
}
pEnum->Release();
}
else
PrintError("Failed to enumerate work items: 0x%X", hr);
return hr;
}
HRESULT CEssTest::DeleteReferences(LPCWSTR szPath)
{
HRESULT hr;
WCHAR szQuery[MAX_PATH * 2];
IEnumWbemClassObject *pEnum = NULL;
// Find the association so we can nuke it.
swprintf(
szQuery,
L"references of {%s}",
szPath);
_bstr_t strWQL = L"WQL",
strQuery = szQuery;
hr =
m_pDefNamespace->ExecQuery(
strWQL,
strQuery,
WBEM_FLAG_FORWARD_ONLY,
NULL,
&pEnum);
if (SUCCEEDED(hr))
{
IWbemClassObjectPtr pObj;
DWORD nCount;
while(SUCCEEDED(hr = pEnum->Next(WBEM_INFINITE, 1, &pObj, &nCount)) &&
nCount > 0)
{
_variant_t vPath;
pObj->Get(L"__PATH", 0, &vPath, NULL, NULL);
hr = m_pDefNamespace->DeleteInstance(
V_BSTR(&vPath), 0, NULL, NULL);
}
pEnum->Release();
}
return hr;
}
HRESULT CEssTest::Init()
{
HRESULT hr;
CoInitializeEx(NULL, COINIT_MULTITHREADED);
if ((hr = CoCreateInstance(
CLSID_WbemLocator,
NULL,
CLSCTX_INPROC_SERVER,
IID_IWbemLocator,
(LPVOID *) &m_pLocator)) == S_OK)
{
hr = GetNamespace(DEF_NAMESPACE, &m_pDefNamespace);
// Get rid of stuff that may have been left around from an aborted
// instance of ESST.
RemoveBindings();
if (SUCCEEDED(hr))
hr = LoadWorkItems();
}
else
PrintError("CoCreateInstance for CLSID_WbemLocator failed: 0x%X\n", hr);
return hr;
}
void CEssTest::RemoveBindings()
{
IEnumWbemClassObject *pEnum = NULL;
_bstr_t strClass;
HRESULT hr;
// Get rid of MSFT_WmiMofConsumer instances.
strClass = L"MSFT_WmiMofConsumer";
hr =
m_pDefNamespace->CreateInstanceEnum(
strClass,
WBEM_FLAG_DEEP | WBEM_FLAG_RETURN_IMMEDIATELY |
WBEM_FLAG_FORWARD_ONLY,
NULL,
&pEnum);
if (SUCCEEDED(hr))
{
IWbemClassObjectPtr pObj;
DWORD nCount;
BOOL bDisplayed = FALSE;
while(SUCCEEDED(hr = pEnum->Next(WBEM_INFINITE, 1, &pObj, &nCount)) &&
nCount > 0)
{
_variant_t vPath;
if (!bDisplayed)
{
PrintStatus("Cleaning up previous consumers...");
bDisplayed = TRUE;
}
pObj->Get(L"__PATH", 0, &vPath, NULL, NULL);
// Get rid of all the bindings using this consumer.
hr = DeleteReferences(V_BSTR(&vPath));
hr = m_pDefNamespace->DeleteInstance(
V_BSTR(&vPath), 0, NULL, NULL);
}
pEnum->Release();
}
else
PrintError("Failed to enumerate MSFT_WmiMofConsumer: 0x%X", hr);
// Get rid of MSFT_EssTestEventFilterToTestFilter instances.
strClass = L"MSFT_EssTestEventFilterToTestFilter";
hr =
m_pDefNamespace->CreateInstanceEnum(
strClass,
WBEM_FLAG_DEEP | WBEM_FLAG_RETURN_IMMEDIATELY |
WBEM_FLAG_FORWARD_ONLY,
NULL,
&pEnum);
if (SUCCEEDED(hr))
{
IWbemClassObjectPtr pObj;
DWORD nCount;
BOOL bDisplayed = FALSE;
while(SUCCEEDED(hr = pEnum->Next(WBEM_INFINITE, 1, &pObj, &nCount)) &&
nCount > 0)
{
_variant_t vPath;
if (!bDisplayed)
{
PrintStatus("Cleaning up previous filters...");
bDisplayed = TRUE;
}
// Nuke the actual __EventFilter.
pObj->Get(L"EventFilter", 0, &vPath, NULL, NULL);
hr = m_pDefNamespace->DeleteInstance(
V_BSTR(&vPath), 0, NULL, NULL);
// Nuke the association.
pObj->Get(L"__PATH", 0, &vPath, NULL, NULL);
hr = m_pDefNamespace->DeleteInstance(
V_BSTR(&vPath), 0, NULL, NULL);
}
pEnum->Release();
}
else
PrintError("Failed to enumerate MSFT_EssTestEventFilterToTestFilter: 0x%X", hr);
}
void CEssTest::Cleanup()
{
for (CWorkItemListIterator i = m_listItems.begin();
i != m_listItems.end();
i++)
{
delete (*i);
}
DeleteCriticalSection(&m_csOutput);
CoUninitialize();
}
void CEssTest::PrintResult(LPSTR szFormat, ...)
{
va_list arglist;
va_start(arglist, szFormat);
Vprintf(LOGLEVEL_RESULTS_ONLY, szFormat, arglist);
}
void CEssTest::PrintError(LPSTR szFormat, ...)
{
va_list arglist;
va_start(arglist, szFormat);
Vprintf(LOGLEVEL_RESULTS_AND_ERRORS, szFormat, arglist);
}
void CEssTest::PrintStatus(LPSTR szFormat, ...)
{
va_list arglist;
va_start(arglist, szFormat);
Vprintf(LOGLEVEL_ALL, szFormat, arglist);
}
void CEssTest::Printf(LOG_LEVEL level, LPSTR szFormat, ...)
{
va_list arglist;
va_start(arglist, szFormat);
Vprintf(level, szFormat, arglist);
}
void CEssTest::Vprintf(LOG_LEVEL level, LPSTR szFormat, va_list arglist)
{
if (level <= m_logLevel)
{
vprintf(szFormat, arglist);
printf("\n");
}
}
HRESULT CEssTest::SpawnInstance(LPCWSTR szClass, IWbemClassObject **ppObj)
{
HRESULT hr;
IWbemClassObject *pClass = NULL;
_bstr_t strClass = szClass;
if (SUCCEEDED(hr = m_pDefNamespace->GetObject(
strClass,
0,
NULL,
&pClass,
NULL)))
{
if (FAILED(hr = pClass->SpawnInstance(
0,
ppObj)))
{
g_essTest.PrintError(
"Unable to spawn an instance of %S: 0x%X", szClass, hr);
}
pClass->Release();
}
else
g_essTest.PrintError(
"Unable to get the %S class: 0x%X", szClass, hr);
return hr;
}
HRESULT CEssTest::GetNamespace(LPCWSTR szNamespace, IWbemServices **ppNamespace)
{
_bstr_t strNamespace = szNamespace;
CNamespaceMapItor item;
HRESULT hr;
wcsupr(strNamespace);
item = m_mapNamespace.find(strNamespace);
if (item != m_mapNamespace.end())
{
*ppNamespace = (*item).second;
(*ppNamespace)->AddRef();
hr = S_OK;
}
else
{
hr =
m_pLocator->ConnectServer(
strNamespace,
NULL, // username
NULL, // password
NULL, // locale
0L, // securityFlags
NULL, // authority (domain for NTLM)
NULL, // context
ppNamespace);
if (SUCCEEDED(hr))
m_mapNamespace[strNamespace] = *ppNamespace;
else
PrintError(
"IWbemLocator::ConnectServer to %S failed: 0x%X\n",
(LPWSTR) strNamespace,
hr);
}
return hr;
}