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