library.cpp

/**
 * @file library.cpp
 * @author Stavros Avramidis
 * @date 5/28/2019
 * @brief definitions of library.h
 */


#include "library.h"

/**
 * @details Convert a wide Unicode String to a UTF8 sring
 * @param wstr
 * @return
 */
std::string utf8_encode(const std::wstring &wstr) {
  if (wstr.empty()) return std::string();

  int size_needed = WideCharToMultiByte(CP_UTF8, 0, &wstr[0], (int) wstr.size(),
                                        nullptr, 0, nullptr, nullptr);
  std::string strTo(size_needed, 0);
  WideCharToMultiByte(CP_UTF8, 0, &wstr[0], (int) wstr.size(), &strTo[0],
                      size_needed, nullptr, nullptr);
  return strTo;
}

/**
 * @details Convert an UTF8 string to a wide Unicode String
 * @param str string to be converted
 * @return a wstring
 */
std::wstring utf8_decode(const std::string &str) {
  if (str.empty()) return std::wstring();
  int size_needed = MultiByteToWideChar(CP_UTF8, 0, &str[0], (int) str.size(), nullptr, 0);
  std::wstring wstrTo(size_needed, 0);
  MultiByteToWideChar(CP_UTF8, 0, &str[0], (int) str.size(), &wstrTo[0], size_needed);
  return wstrTo;
}

Wmi::Wmi() : pLoc(nullptr), pSvc(nullptr) {};

Wmi::~Wmi() {
  pSvc->Release();
  pLoc->Release();
  CoUninitialize();
}

HRESULT Wmi::init() {
  HRESULT hres;

  // Initialize COM. ------------------------------------------
  hres = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
  if (FAILED(hres)) {
    return hres;
  }

  // Set general COM security levels --------------------------
  hres = CoInitializeSecurity(
      nullptr,
      -1,                          // COM authentication
      nullptr,                     // Authentication services
      nullptr,                     // Reserved
      RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication
      RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation
      nullptr,                     // Authentication info
      EOAC_NONE,                   // Additional capabilities
      nullptr                      // Reserved
  );

  if (FAILED(hres)) {
    CoUninitialize();
    return hres;
  }

  // Obtain the initial locator to WMI -------------------------
  hres = CoCreateInstance(
      CLSID_WbemLocator,
      nullptr,
      CLSCTX_INPROC_SERVER,
      IID_IWbemLocator, (LPVOID *) &pLoc);

  if (FAILED(hres)) {
    CoUninitialize();
    return hres;
  }

  hres = pLoc->ConnectServer(
      _bstr_t(L"ROOT\\CIMV2"), // Object path of WMI namespace
      nullptr,                    // User name. NULL = current user
      nullptr,                    // User password. NULL = current
      nullptr,                 // Locale. NULL indicates current
      0,                    // Security flags.
      nullptr,                 // Authority (for example, Kerberos)
      nullptr,                 // Context object
      &pSvc                    // pointer to IWbemServices proxy
  );

  if (FAILED(hres)) {
    pLoc->Release();
    CoUninitialize();
    return hres;
  }

  // Set security levels on the proxy -------------------------
  hres = CoSetProxyBlanket(
      pSvc,                        // Indicates the proxy to set
      RPC_C_AUTHN_WINNT,           // RPC_C_AUTHN_xxx
      RPC_C_AUTHZ_NONE,            // RPC_C_AUTHZ_xxx
      nullptr,                        // Server principal name
      RPC_C_AUTHN_LEVEL_CALL,      // RPC_C_AUTHN_LEVEL_xxx
      RPC_C_IMP_LEVEL_IMPERSONATE, // RPC_C_IMP_LEVEL_xxx
      nullptr,                        // client identity
      EOAC_NONE                    // proxy capabilities
  );

  if (FAILED(hres)) {
    pSvc->Release();
    pLoc->Release();
    CoUninitialize();
    return hres;
  }
  // Initialised WMI successfully
  return S_OK;
}

HRESULT Wmi::query(const std::string &queryStr,
                   std::vector<QueryObj> &queryVectorOut,
                   const AdditionalFilters *filters) {
  HRESULT hres;

  IEnumWbemClassObject *pEnumerator = nullptr;

  // Make the WMI query
  hres = pSvc->ExecQuery(
      bstr_t("WQL"),
      bstr_t(queryStr.c_str()),
      WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
      nullptr,
      &pEnumerator);

  if (FAILED(hres))
    return hres;

  IWbemClassObject *pclsObj = nullptr;
  ULONG uReturn = 0;

  while (pEnumerator) {
    pEnumerator->Next(WMI_WAIT_TIME, 1, &pclsObj, &uReturn);

    if (0==uReturn)
      break;

    VARIANT vtProp;

    if (filters) {
      for (auto &filter: *filters) {
        hres = pclsObj->Get(utf8_decode(filter.first).c_str(), 0, &vtProp, nullptr, nullptr);
        if (FAILED(hres)) continue;
        auto val = utf8_encode(vtProp.bstrVal);
        if (!std::regex_match(val, filter.second)) {
          goto _NextElement;
        }
      }
    }

    SAFEARRAY *sfArray;
    LONG lstart, lend;


    //Get Wmi objects names
    hres = pclsObj->GetNames(nullptr, WBEM_FLAG_ALWAYS, nullptr, &sfArray);
    if (FAILED(hres))
      continue;

    // Find safe array boundaries
    SafeArrayGetLBound(sfArray, 1, &lstart);
    SafeArrayGetUBound(sfArray, 1, &lend);

    BSTR *pbstr;
    hres = SafeArrayAccessData(sfArray, (void HUGEP **) &pbstr);
    LONG nIdx = 0;

    if (FAILED(hres)) continue;

    else {
      CIMTYPE pType;
      QueryObj item;
      for (nIdx = lstart; nIdx < lend; nIdx++) {
        hres = pclsObj->Get(pbstr[nIdx], 0, &vtProp, &pType, 0);

        if (FAILED(hres)) continue;

        if (vtProp.vt==VT_NULL) continue;

        if ((pType==CIM_STRING || pType==CIM_REFERENCE) && pType!=CIM_EMPTY && pType!=CIM_ILLEGAL)
          item.emplace(utf8_encode(pbstr[nIdx]), utf8_encode(vtProp.bstrVal));

        VariantClear(&vtProp);
      }

      hres = pclsObj->Get(L"Dependent", 0, &vtProp, &pType, nullptr);
      if (pType!=CIM_EMPTY && pType!=CIM_ILLEGAL && SUCCEEDED(hres))
        item.emplace(utf8_encode(pbstr[nIdx]), utf8_encode(vtProp.bstrVal));

      // Push item to vector
      queryVectorOut.emplace_back(item);

      // Empty sfArray
      SafeArrayUnaccessData(sfArray);

      SafeArrayDestroy(sfArray);  // Delete sfArray
      sfArray = nullptr;          // Avoid dangling pointers
    }

    _NextElement:
    VariantClear(&vtProp);      // Clear vtProp
    pclsObj->Release();         // Release pclsObj
  }

  pEnumerator->Release();
  pEnumerator = nullptr;

  return S_OK;
}