Update fundamental WMI classes

5 files changed, 33 insertions, 31 deletions

diff --git a/windns/src/windns/wmi/connection.cpp b/windns/src/windns/wmi/connection.cpp
index a2588b139e..225586a11d 100644
--- a/windns/src/windns/wmi/connection.cpp
+++ b/windns/src/windns/wmi/connection.cpp
@@ -1,6 +1,6 @@
 #include "stdafx.h"
 #include "connection.h"
-#include "windns/comlol.h"
+#include "windns/comhelpers.h"
 #include <stdexcept>
 #define _WIN32_DCOM
 #include <windows.h>
@@ -15,6 +15,7 @@ const wchar_t *LiteralNamespace(wmi::Connection::Namespace ns)
 	{
 		case wmi::Connection::Namespace::Default: return L"root\\Default";
 		case wmi::Connection::Namespace::Cimv2: return L"root\\CIMV2";
+		case wmi::Connection::Namespace::StandardCimv2: return L"root\\StandardCIMV2";
 		default:
 		{
 			throw std::logic_error("Missing case handler in switch clause");
@@ -32,6 +33,14 @@ Connection::Connection(Namespace ns) : m_queryLanguage(L"WQL")
 	auto status = CoCreateInstance(CLSID_WbemLocator, nullptr, CLSCTX_INPROC_SERVER,
 		IID_IWbemLocator, (LPVOID *)&m_locator);
 
+	if (CO_E_NOTINITIALIZED == status)
+	{
+		VALIDATE_COM(CoInitializeEx(nullptr, COINIT_MULTITHREADED), "Initialize COM");
+
+		status = CoCreateInstance(CLSID_WbemLocator, nullptr, CLSCTX_INPROC_SERVER,
+			IID_IWbemLocator, (LPVOID *)&m_locator);
+	}
+
 	VALIDATE_COM(status, "Create COM locator instance");
 
 	status = m_locator->ConnectServer(_bstr_t(LiteralNamespace(ns)), nullptr, nullptr,
@@ -45,7 +54,7 @@ Connection::Connection(Namespace ns) : m_queryLanguage(L"WQL")
 	VALIDATE_COM(status, "Configure COM services auth");
 }
 
-ResultSet Connection::Query(const wchar_t *query)
+ResultSet Connection::query(const wchar_t *query)
 {
 	CComPtr<IEnumWbemClassObject> result;
 
diff --git a/windns/src/windns/wmi/connection.h b/windns/src/windns/wmi/connection.h
index 37d658bdaf..07e5066d1c 100644
--- a/windns/src/windns/wmi/connection.h
+++ b/windns/src/windns/wmi/connection.h
@@ -6,6 +6,9 @@
 #include <windows.h>
 #include <atlbase.h>
 #include <comutil.h>
+#include <comdef.h>
+
+#pragma comment(lib, "wbemuuid.lib")
 
 namespace wmi
 {
@@ -17,24 +20,21 @@ public:
 	enum class Namespace
 	{
 		Default,
-		Cimv2
+		Cimv2,
+		StandardCimv2
 	};
 
 	explicit Connection(Namespace ns);
 
-	ResultSet Query(const wchar_t *query) override;
+	ResultSet query(const wchar_t *query) override;
 
-	// TODO: Move to shared base class.
-	ResultSet Query(const std::wstring &str)
+	CComPtr<IWbemServices> services() override
 	{
-		return Query(str.c_str());
+		return m_services;
 	}
 
 private:
 
-	Connection(const Connection &) = delete;
-	Connection &operator=(const Connection &) = delete;
-
 	CComPtr<IWbemLocator> m_locator;
 	CComPtr<IWbemServices> m_services;
 
diff --git a/windns/src/windns/wmi/iconnection.h b/windns/src/windns/wmi/iconnection.h
index 58b5841b91..8683dd6ac5 100644
--- a/windns/src/windns/wmi/iconnection.h
+++ b/windns/src/windns/wmi/iconnection.h
@@ -1,6 +1,8 @@
 #pragma once
 
 #include "resultset.h"
+#include <atlbase.h>
+#include <wbemidl.h>
 
 namespace wmi
 {
@@ -11,7 +13,8 @@ struct IConnection
 	{
 	}
 
-	virtual ResultSet Query(const wchar_t *query) = 0;
+	virtual ResultSet query(const wchar_t *query) = 0;
+	virtual CComPtr<IWbemServices> services() = 0;
 };
 
 }
diff --git a/windns/src/windns/wmi/resultset.cpp b/windns/src/windns/wmi/resultset.cpp
index b5914ff6d2..cea147fa76 100644
--- a/windns/src/windns/wmi/resultset.cpp
+++ b/windns/src/windns/wmi/resultset.cpp
@@ -1,6 +1,6 @@
 #include "stdafx.h"
 #include "resultset.h"
-#include "windns/comlol.h"
+#include "windns/comhelpers.h"
 #include "libcommon/error.h"
 
 namespace wmi
@@ -12,6 +12,11 @@ ResultSet::ResultSet(CComPtr<IEnumWbemClassObject> rs) : m_resultset(rs)
 
 bool ResultSet::advance()
 {
+	if (nullptr != m_result)
+	{
+		m_result.Release();
+	}
+
 	ULONG dummy;
 
 	const auto status = m_resultset->Next(WBEM_INFINITE, 1, &m_result, &dummy);
@@ -21,18 +26,7 @@ bool ResultSet::advance()
 	return WBEM_S_FALSE != status;
 }
 
-_variant_t ResultSet::getResultProperty(const std::wstring &name)
-{
-	_variant_t val;
-
-	auto status = m_result->Get(name.c_str(), 0, &val, nullptr, nullptr);
-
-	VALIDATE_COM(status, "Retrieve COM property value");
-
-	return val;
-}
-
-CComPtr<IWbemClassObject> ResultSet::getResult()
+CComPtr<IWbemClassObject> ResultSet::result()
 {
 	return m_result;
 }
diff --git a/windns/src/windns/wmi/resultset.h b/windns/src/windns/wmi/resultset.h
index bcdefa6c3c..6139a6051a 100644
--- a/windns/src/windns/wmi/resultset.h
+++ b/windns/src/windns/wmi/resultset.h
@@ -3,7 +3,6 @@
 #include <string>
 #include <atlbase.h>
 #include <wbemidl.h>
-#include <comutil.h>
 
 namespace wmi
 {
@@ -14,20 +13,17 @@ public:
 
 	ResultSet(CComPtr<IEnumWbemClassObject> rs);
 
+	ResultSet(const ResultSet &) = delete;
+	ResultSet &operator=(const ResultSet &) = delete;
 	ResultSet(ResultSet &&) = default;
 	ResultSet &operator=(ResultSet &&) = default;
 
 	bool advance();
 
-	_variant_t getResultProperty(const std::wstring &name);
-
-	CComPtr<IWbemClassObject> getResult();
+	CComPtr<IWbemClassObject> result();
 
 private:
 
-	ResultSet(const ResultSet &) = delete;
-	ResultSet &operator=(const ResultSet &) = delete;
-
 	CComPtr<IEnumWbemClassObject> m_resultset;
 	CComPtr<IWbemClassObject> m_result;
 };