os_win32/wmiquery.cpp

   1 /*
   2  * os_win32/wmiquery.cpp
   3  *
   4  * Home page of code is: http://www.smartmontools.org
   5  *
   6  * Copyright (C) 2011-13 Christian Franke
   7  *
   8  * SPDX-License-Identifier: GPL-2.0-or-later
   9  */
  10
  11 #include "config.h"
  12 #define WINVER 0x0400
  13 #define _WIN32_WINNT WINVER
  14
  15 #include "wmiquery.h"
  16
  17 #include <stdio.h>
  18
  19 const char * wmiquery_cpp_cvsid = "$Id: wmiquery.cpp 4760 2018-08-19 18:45:53Z chrfranke $"
  20   WMIQUERY_H_CVSID;
  21
  22
  23 /////////////////////////////////////////////////////////////////////////////
  24 // com_bstr
  25
  26 com_bstr::com_bstr(const char * str)
  27 : m_bstr(0)
  28 {
  29   int sz = MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, str, -1, (LPWSTR)0, 0);
  30   if (sz <= 0)
  31     return;
  32   m_bstr = SysAllocStringLen((OLECHAR*)0, sz-1);
  33   if (!m_bstr)
  34     return; // throw std::bad_alloc
  35   MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, str, -1, m_bstr, sz);
  36 }
  37
  38 bool com_bstr::to_str(const BSTR & bstr, std::string & str)
  39 {
  40   if (!bstr)
  41     return false;
  42   int sz = WideCharToMultiByte(CP_ACP, 0, bstr, -1, (LPSTR)0, 0, (LPCSTR)0, (LPBOOL)0);
  43   if (sz <= 0)
  44     return false;
  45   char * buf = new char[sz];
  46   WideCharToMultiByte(CP_ACP, 0, bstr, -1, buf, sz, (LPCSTR)0, (LPBOOL)0);
  47   str = buf;
  48   delete [] buf;
  49   return true;
  50 }
  51
  52
  53 /////////////////////////////////////////////////////////////////////////////
  54 // wbem_object
  55
  56 std::string wbem_object::get_str(const char * name) /*const*/
  57 {
  58   std::string s;
  59   if (!m_intf)
  60     return s;
  61
  62   VARIANT var; VariantInit(&var);
  63   if (m_intf->Get(com_bstr(name), 0L, &var, (CIMTYPE*)0, (LPLONG)0) /* != WBEM_S_NO_ERROR */)
  64     return s;
  65
  66   if (var.vt == VT_BSTR)
  67     com_bstr::to_str(var.bstrVal, s);
  68   VariantClear(&var);
  69   return s;
  70 }
  71
  72
  73 /////////////////////////////////////////////////////////////////////////////
  74 // wbem_enumerator
  75
  76 bool wbem_enumerator::next(wbem_object & obj)
  77 {
  78   if (!m_intf)
  79     return false;
  80
  81   ULONG n = 0;
  82   HRESULT rc = m_intf->Next(5000 /*5s*/, 1 /*count*/, obj.m_intf.replace(), &n);
  83   if (FAILED(rc) || n != 1)
  84     return false;
  85   return true;
  86 }
  87
  88
  89 /////////////////////////////////////////////////////////////////////////////
  90 // wbem_services
  91
  92 const CLSID xCLSID_WbemLocator = {0x4590f811, 0x1d3a, 0x11d0, {0x89, 0x1f, 0x00, 0xaa, 0x00, 0x4b, 0x2e, 0x24}};
  93 const IID   xIID_IWbemLocator  = {0xdc12a687, 0x737f, 0x11cf, {0x88, 0x4d, 0x00, 0xaa, 0x00, 0x4b, 0x2e, 0x24}};
  94
  95 bool wbem_services::connect()
  96 {
  97   // Init COM during first call.
  98   static HRESULT init_rc = -1;
  99   static bool init_tried = false;
 100   if (!init_tried) {
 101     init_tried = true;
 102     init_rc = CoInitialize((LPVOID)0);
 103   }
 104   if (!(init_rc == S_OK  || init_rc == S_FALSE))
 105     return false;
 106
 107   /// Create locator.
 108   com_intf_ptr<IWbemLocator> locator;
 109   HRESULT rc = CoCreateInstance(xCLSID_WbemLocator, (LPUNKNOWN)0,
 110     CLSCTX_INPROC_SERVER, xIID_IWbemLocator, (LPVOID*)locator.replace());
 111   if (FAILED(rc))
 112     return false;
 113
 114   // Set timeout flag if supported.
 115   long flags = 0;
 116   OSVERSIONINFOA ver; ver.dwOSVersionInfoSize = sizeof(ver);
 117   if (GetVersionExA(&ver) && ver.dwPlatformId == VER_PLATFORM_WIN32_NT
 118     && (    ver.dwMajorVersion >= 6 // Vista
 119         || (ver.dwMajorVersion == 5 && ver.dwMinorVersion >= 1))) // XP
 120     flags = WBEM_FLAG_CONNECT_USE_MAX_WAIT; // return in 2min or less
 121
 122   // Connect to local server.
 123   rc = locator->ConnectServer(com_bstr("\\\\.\\root\\cimv2"),
 124     (BSTR)0, (BSTR)0, (BSTR)0, // User, Password, Locale
 125     flags, (BSTR)0, (IWbemContext*)0, m_intf.replace());
 126   if (FAILED(rc))
 127     return false;
 128
 129   // Set authentication information,
 130   rc = CoSetProxyBlanket(m_intf.get(), RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE,
 131     (OLECHAR*)0, RPC_C_AUTHN_LEVEL_CALL, RPC_C_IMP_LEVEL_IMPERSONATE,
 132     (RPC_AUTH_IDENTITY_HANDLE*)0, EOAC_NONE);
 133   if (FAILED(rc)) {
 134     m_intf.reset();
 135     return false;
 136   }
 137
 138   return true;
 139 }
 140
 141 bool wbem_services::vquery(wbem_enumerator & result, const char * qstr, va_list args) /*const*/
 142 {
 143   if (!m_intf)
 144     return false;
 145
 146   char qline[1024];
 147   vsnprintf(qline, sizeof(qline), qstr, args);
 148   qline[sizeof(qline)-1] = 0;
 149
 150   HRESULT rc = m_intf->ExecQuery(
 151     com_bstr("WQL"), com_bstr(qline),
 152     WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
 153     (IWbemContext*)0, result.m_intf.replace());
 154   if (FAILED(rc))
 155     return false;
 156
 157   return true;
 158 }
 159
 160 bool wbem_services::vquery1(wbem_object & obj, const char * qstr, va_list args) /*const*/
 161 {
 162   wbem_enumerator result;
 163   if (!vquery(result, qstr, args))
 164     return false;
 165
 166   if (!result.next(obj))
 167     return false;
 168
 169   wbem_object peek;
 170   if (result.next(peek))
 171     return false;
 172
 173   return true;
 174 }
 175
 176 bool wbem_services::query(wbem_enumerator & result, const char * qstr, ...) /*const*/
 177 {
 178   va_list args; va_start(args, qstr);
 179   bool ok = vquery(result, qstr, args);
 180   va_end(args);
 181   return ok;
 182 }
 183
 184 bool wbem_services::query1(wbem_object & obj, const char * qstr, ...) /*const*/
 185 {
 186   va_list args; va_start(args, qstr);
 187   bool ok = vquery1(obj, qstr, args);
 188   va_end(args);
 189   return ok;
 190 }