CCEngineServer/CMakeProject1/SystemParamsProvider_win.cpp
2021-08-29 23:21:55 +02:00

421 lines
14 KiB
C++

#include "SystemParamsProvider_win.h"
struct IUnknown; // Workaround for "combaseapi.h(229): error C2187: syntax error: 'identifier' was unexpected here" when using /permissive-
#define _WIN32_DCOM
#include <comdef.h>
#include <Wbemidl.h>
#include <boost/algorithm/string/join.hpp>
#include <future>
#include <iostream>
#include <sstream>
#pragma comment(lib, "wbemuuid.lib")
namespace
{
std::string convertWCSToMBS(const wchar_t* pstr, long wslen)
{
int len = ::WideCharToMultiByte(CP_ACP, 0, pstr, wslen, NULL, 0, NULL, NULL);
std::string dblstr(len, '\0');
len = ::WideCharToMultiByte(CP_ACP, 0 /* no flags */,
pstr, wslen /* not necessary NULL-terminated */,
&dblstr[0], len,
NULL, NULL /* no default char */);
return dblstr;
}
std::string convertBSTRToMBS(BSTR bstr)
{
int wslen = ::SysStringLen(bstr);
return convertWCSToMBS((wchar_t*)bstr, wslen);
}
BSTR convertMBSToBSTR(const std::string& str)
{
int wslen = ::MultiByteToWideChar(CP_ACP, 0 /* no flags */,
str.data(), static_cast<int>(str.length()),
NULL, 0);
BSTR wsdata = ::SysAllocStringLen(NULL, wslen);
::MultiByteToWideChar(CP_ACP, 0 /* no flags */,
str.data(), static_cast<int>(str.length()),
wsdata, wslen);
return wsdata;
}
class QuerySink final : public IWbemObjectSink
{
public:
typedef std::vector<std::string> Fields_type;
typedef std::vector<std::vector<std::string>> Result_type;
public:
QuerySink(const std::string& wmiClass, const Fields_type& fields)
: m_wmiClass(wmiClass)
, m_fields(fields)
{
m_lRef = 0;
m_bDone = false;
InitializeCriticalSection(&m_threadLock);
}
~QuerySink() {
m_bDone = true;
DeleteCriticalSection(&m_threadLock);
}
public:
virtual ULONG STDMETHODCALLTYPE AddRef() override
{
return InterlockedIncrement(&m_lRef);
}
virtual ULONG STDMETHODCALLTYPE Release() override
{
LONG lRef = InterlockedDecrement(&m_lRef);
if (lRef == 0)
delete this;
return lRef;
}
virtual HRESULT STDMETHODCALLTYPE QueryInterface(REFIID riid, void** ppv) override
{
if (riid == IID_IUnknown || riid == IID_IWbemObjectSink)
{
*ppv = (IWbemObjectSink *)this;
AddRef();
return WBEM_S_NO_ERROR;
}
else return E_NOINTERFACE;
}
virtual HRESULT STDMETHODCALLTYPE Indicate(
LONG lObjectCount,
IWbemClassObject __RPC_FAR *__RPC_FAR *apObjArray
) override
{
for (int i = 0; i < lObjectCount; i++)
{
std::vector<std::string> lineResult;
lineResult.reserve(m_fields.size());
for (const auto& field : m_fields)
{
HRESULT hres = S_OK;
VARIANT varName;
hres = apObjArray[i]->Get(convertMBSToBSTR(field),
0, &varName, 0, 0);
if (FAILED(hres))
{
std::ostringstream ostr;
ostr << "Failed to get the data from the query. Error code = 0x"
<< std::hex << hres;
std::cerr << ostr.str() << std::endl;
return WBEM_E_FAILED; // Program has failed.
}
lineResult.push_back(convertBSTRToMBS(V_BSTR(&varName)));
}
m_result.push_back(lineResult);
}
return WBEM_S_NO_ERROR;
}
virtual HRESULT STDMETHODCALLTYPE SetStatus(
/* [in] */ LONG lFlags,
/* [in] */ HRESULT hResult,
/* [in] */ BSTR strParam,
/* [in] */ IWbemClassObject __RPC_FAR *pObjParam
) override
{
if (lFlags == WBEM_STATUS_COMPLETE)
{
// Call complete
EnterCriticalSection(&m_threadLock);
m_bDone = true;
LeaveCriticalSection(&m_threadLock);
m_resultPromise.set_value(std::move(m_result));
}
else if (lFlags == WBEM_STATUS_PROGRESS)
{
// Call in progress...
}
return WBEM_S_NO_ERROR;
}
std::future<Result_type> getFuture()
{
return m_resultPromise.get_future();
}
bool IsDone() const
{
bool done = true;
EnterCriticalSection(&m_threadLock);
done = m_bDone;
LeaveCriticalSection(&m_threadLock);
return done;
}
private:
LONG m_lRef;
bool m_bDone;
mutable CRITICAL_SECTION m_threadLock; // for thread safety
const std::string m_wmiClass;
const Fields_type m_fields;
std::vector<std::vector<std::string>> m_result;
std::promise<Result_type> m_resultPromise;
};
void initializeCOM()
{
// We need to initialize COM for SystemParamsProvider
HRESULT hres = 0;
// Initialize COM. ------------------------------------------
hres = CoInitializeEx(0, COINIT_MULTITHREADED);
if (FAILED(hres))
{
std::ostringstream ostr;
ostr << "Failed to initialize COM library. Error code = 0x" << std::hex << hres;
throw std::runtime_error(ostr.str());
}
// Set general COM security levels --------------------------
hres = CoInitializeSecurity(NULL,
-1, // COM authentication
NULL, // Authentication services
NULL, // Reserved
RPC_C_AUTHN_LEVEL_DEFAULT, // Default authentication
RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation
NULL, // Authentication info
EOAC_NONE, // Additional capabilities
NULL); // Reserved
if (FAILED(hres) && hres != RPC_E_TOO_LATE) // Ignore error if the security has already been initialized
{
CoUninitialize();
std::ostringstream ostr;
ostr << "Failed to initialize security. Error code = 0x" << std::hex << hres;
throw std::runtime_error(ostr.str());
}
}
}
namespace detail
{
class SystemParamsProvider_winImpl final
{
public:
SystemParamsProvider_winImpl()
{
HRESULT hres = 0;
initializeCOM();
// Obtain the initial locator to WMI -------------------------
hres = CoCreateInstance(
CLSID_WbemLocator,
0,
CLSCTX_INPROC_SERVER,
IID_IWbemLocator, (LPVOID *)&m_pLoc);
if (FAILED(hres))
{
CoUninitialize();
std::ostringstream ostr;
ostr << "Failed to create IWbemLocator object. Err code = 0x" << std::hex << hres;
throw std::runtime_error(ostr.str());
}
// Connect to WMI through the IWbemLocator::ConnectServer method
// Connect to the local root\cimv2 namespace
// and obtain pointer pSvc to make IWbemServices calls.
hres = m_pLoc->ConnectServer(_bstr_t(L"ROOT\\CIMV2"),
NULL,
NULL,
0,
NULL,
0,
0,
&m_pSvc);
if (FAILED(hres))
{
m_pLoc->Release();
CoUninitialize();
std::ostringstream ostr;
ostr << "Could not connect. Error code = 0x" << std::hex << hres;
throw std::runtime_error(ostr.str());
}
// Set security levels on the proxy -------------------------
hres = CoSetProxyBlanket(m_pSvc, // Indicates the proxy to set
RPC_C_AUTHN_WINNT, // RPC_C_AUTHN_xxx
RPC_C_AUTHZ_NONE, // RPC_C_AUTHZ_xxx
NULL, // Server principal name
RPC_C_AUTHN_LEVEL_CALL, // RPC_C_AUTHN_LEVEL_xxx
RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
NULL, // client identity
EOAC_NONE); // proxy capabilities
if (FAILED(hres))
{
m_pSvc->Release();
m_pLoc->Release();
CoUninitialize();
std::ostringstream ostr;
ostr << "Could not set proxy blanket. Error code = 0x" << std::hex << hres;
throw std::runtime_error(ostr.str());
}
}
~SystemParamsProvider_winImpl()
{
m_pSvc->Release();
m_pLoc->Release();
CoUninitialize();
}
public:
SystemParams retrieveSystemParams()
{
auto biosFuture = queryWmiInfo("Win32_BIOS", { "SerialNumber" });
auto computerSystemProductFuture = queryWmiInfo("Win32_ComputerSystemProduct", { "UUID" });
auto diskDriveFuture = queryWmiInfo("Win32_DiskDrive", { "SerialNumber" }, "DeviceId = \"\\\\\\\\.\\\\PHYSICALDRIVE0\"");
auto osFuture = queryWmiInfo("Win32_OperatingSystem", { "SerialNumber" });
auto nicFuture = queryWmiInfo("Win32_NetworkAdapter", { "MACAddress" }, "PhysicalAdapter = TRUE");
auto mainboardFuture = queryWmiInfo("Win32_BaseBoard", { "SerialNumber" });
SystemParams result;
result[SystemParamTypes::biosSerialNum] = processBiosData(biosFuture.get());
result[SystemParamTypes::computerUUID] = processComputerSystemProductData(computerSystemProductFuture.get());
result[SystemParamTypes::diskSerialNum] = processDiskDrive(diskDriveFuture.get());
result[SystemParamTypes::osId] = processOs(osFuture.get());
result[SystemParamTypes::nicMac] = processNic(nicFuture.get());
result[SystemParamTypes::mainboardSerialNum] = processMainBoard(mainboardFuture.get());
return result;
}
private:
typedef std::vector<std::vector<std::string>> QueryResult_type;
private:
std::future<QueryResult_type> queryWmiInfo(const std::string& wmiClass, const std::vector<std::string>& fieldNames, const std::string& condition = "")
{
HRESULT hres = 0;
std::ostringstream queryStr;
queryStr << "SELECT " << boost::algorithm::join(fieldNames, ",") << " FROM " << wmiClass;
if (!condition.empty())
{
queryStr << " WHERE " << condition;
}
// Use the IWbemServices pointer to make requests of WMI ----
QuerySink* pResponseSink = new QuerySink(wmiClass, fieldNames);
pResponseSink->AddRef();
auto resultFuture = pResponseSink->getFuture();
hres = m_pSvc->ExecQueryAsync(bstr_t("WQL"),
convertMBSToBSTR(queryStr.str()),
WBEM_FLAG_BIDIRECTIONAL,
NULL,
pResponseSink);
if (FAILED(hres))
{
pResponseSink->Release();
std::ostringstream ostr;
ostr << "Query for " << boost::algorithm::join(fieldNames, ", ") << " failed. Error code = 0x" << std::hex << hres;
throw std::runtime_error(ostr.str());
}
return resultFuture;
}
static std::string processBiosData(const QueryResult_type& biosInfo)
{
if (biosInfo.size() > 0)
{
return biosInfo[0][0];
}
return {};
}
static std::string processComputerSystemProductData(const QueryResult_type& csInfo)
{
if (csInfo.size() > 0)
{
return csInfo[0][0];
}
return {};
}
static std::string processDiskDrive(const QueryResult_type& csInfo)
{
if (csInfo.size() > 0)
{
return csInfo[0][0];
}
return {};
}
static std::string processOs(const QueryResult_type& csInfo)
{
if (csInfo.size() > 0)
{
return csInfo[0][0];
}
return {};
}
static std::string processNic(const QueryResult_type& csInfo)
{
if (csInfo.size() > 0)
{
return csInfo[0][0];
}
return {};
}
static std::string processMainBoard(const QueryResult_type& csInfo)
{
if (csInfo.size() > 0)
{
return csInfo[0][0];
}
return {};
}
private:
IWbemLocator* m_pLoc = nullptr;
IWbemServices *m_pSvc = nullptr;
};
}
SystemParamsProvider_win::SystemParamsProvider_win()
: m_impl(std::make_unique<detail::SystemParamsProvider_winImpl>())
{
}
SystemParamsProvider_win::~SystemParamsProvider_win()
{
}
SystemParams SystemParamsProvider_win::retrieveSystemParams()
{
return m_impl->retrieveSystemParams();
}