source: nscp/modules/CheckWMI/WMIQuery.cpp

Last change on this file was 94bb8cc, checked in by Michael Medin <michael@…>, 7 weeks ago

Fixed issue with commands escaping
Changed a lot of the utf16 into utf8 as well as major API refactoring
  • Property mode set to 100644
File size: 12.8 KB
Line 
1/**************************************************************************
2*   Copyright (C) 2004-2007 by Michael Medin <michael@medin.name>         *
3*                                                                         *
4*   This code is part of NSClient++ - http://trac.nakednuns.org/nscp      *
5*                                                                         *
6*   This program is free software; you can redistribute it and/or modify  *
7*   it under the terms of the GNU General Public License as published by  *
8*   the Free Software Foundation; either version 2 of the License, or     *
9*   (at your option) any later version.                                   *
10*                                                                         *
11*   This program is distributed in the hope that it will be useful,       *
12*   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
13*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
14*   GNU General Public License for more details.                          *
15*                                                                         *
16*   You should have received a copy of the GNU General Public License     *
17*   along with this program; if not, write to the                         *
18*   Free Software Foundation, Inc.,                                       *
19*   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
20***************************************************************************/
21#include "StdAfx.h"
22#include ".\wmiquery.h"
23
24#include <map>
25
26#include <boost/optional.hpp>
27
28#include <objidl.h>
29#include <Wbemidl.h>
30#include <WMIUtils.h>
31
32#include <error_com.hpp>
33
34std::wstring WMIQuery::sanitize_string(LPTSTR in) {
35        wchar_t *p = in;
36        while (*p) {
37                if (p[0] < ' ' || p[0] > '}')
38                        p[0] = '.';
39                p++;
40        }
41        return in;
42}
43
44struct identidy_container {
45        identidy_container(std::wstring domain, std::wstring username, std::wstring password)
46                : domain(domain)
47                , username(username)
48                , password(password)
49        {
50                setup();
51        }
52        identidy_container(const identidy_container &other)
53                : domain(other.domain)
54                , username(other.username)
55                , password(other.password)
56        {
57                setup();
58        }
59        const identidy_container& operator=(const identidy_container &other) {
60                domain = other.domain;
61                username = other.username;
62                password = other.password;
63                setup();
64                return *this;
65        }
66
67        COAUTHIDENTITY auth_identity;
68        std::wstring domain;
69        std::wstring username;
70        std::wstring password;
71
72private:
73        void setup() {
74                memset(&auth_identity, 0, sizeof(COAUTHIDENTITY));
75                auth_identity.PasswordLength = password.size();
76                auth_identity.Password = (USHORT*)password.c_str();
77                auth_identity.UserLength = username.size();
78                auth_identity.User = (USHORT*)username.c_str();
79                auth_identity.DomainLength = domain.size();
80                auth_identity.Domain = (USHORT*)domain.c_str();
81                auth_identity.Flags = SEC_WINNT_AUTH_IDENTITY_UNICODE;
82        }
83
84};
85
86identidy_container get_identity(const std::wstring &username, const std::wstring &password) {
87        std::wstring::size_type pos = username.find('\\');
88        if(pos == std::string::npos) {
89                return identidy_container(_T(""), username, password);
90        }
91        return identidy_container(username.substr(0, pos), username.substr(pos+1), password);
92}
93
94void set_proxy_blanket(IUnknown *pProxy, const std::wstring &user, const std::wstring &password) {
95        if (user.empty() || password.empty())
96                return;
97        identidy_container auth = get_identity(user, password);
98        HRESULT hr = CoSetProxyBlanket(pProxy, RPC_C_AUTHN_DEFAULT, RPC_C_AUTHZ_DEFAULT, COLE_DEFAULT_PRINCIPAL, RPC_C_AUTHN_LEVEL_DEFAULT,
99                RPC_C_IMP_LEVEL_IMPERSONATE, &auth.auth_identity, EOAC_NONE );
100        if (FAILED(hr))
101                throw WMIException("CoSetProxyBlanket failed: " + ComError::getComError(ComError::getWMIError(hr)));
102}
103
104
105CComPtr<IWbemServices> create_service(std::wstring ns, std::wstring user, std::wstring password) {
106        CComPtr<IWbemLocator> locator;
107        HRESULT hr = CoCreateInstance(CLSID_WbemLocator, NULL, CLSCTX_INPROC_SERVER, IID_IWbemLocator, reinterpret_cast< void** >(&locator));
108        if (FAILED(hr))
109                throw WMIException("CoCreateInstance for CLSID_WbemAdministrativeLocator failed!", hr);
110
111        CComBSTR strNamespace(ns.c_str());
112        CComBSTR strUser(user.c_str());
113        CComBSTR strPassword(password.c_str());
114
115        CComPtr< IWbemServices > service;
116        // @todo: WBEM_FLAG_CONNECT_USE_MAX_WAIT
117        hr = locator->ConnectServer(strNamespace, user.empty()?NULL:strUser, password.empty()?NULL:strPassword, NULL, NULL, 0, NULL, &service );
118        if (FAILED(hr))
119                throw WMIException("ConnectServer failed: namespace=" + utf8::cvt<std::string>(ns) + ", user=" + utf8::cvt<std::string>(user), hr);
120
121        set_proxy_blanket(service, user, password);
122
123        return service;
124}
125
126template <class T, enum VARENUM U>
127std::wstring parse_item(T value);
128
129template<>
130std::wstring parse_item<SHORT,VT_BOOL>(SHORT value) {
131        return value?_T("TRUE"):_T("FALSE");
132}
133template<>
134std::wstring parse_item<BSTR,VT_BSTR>(BSTR value) {
135        return OLE2T(value);
136}
137template<>
138std::wstring parse_item<UINT,VT_UI1>(UINT value) {
139        return strEx::itos(value);
140}
141template<>
142std::wstring parse_item<LONG,VT_I4>(LONG value) {
143        return strEx::itos(value);
144}
145template<class T, enum VARENUM U>
146std::wstring array_to_string(std::wstring tag, CComVariant &vValue) {
147        SAFEARRAY* paArray = vValue.parray;
148
149        T * array = NULL;
150
151        SafeArrayAccessData(paArray, (void**)&array);
152
153        long lLBound = 0;
154        long lUBound = 0;
155        long nCount = 0;
156
157        if (FAILED(SafeArrayGetLBound(paArray, 1, &lLBound)) ||
158                FAILED(SafeArrayGetUBound(paArray, 1, &lUBound)))
159        {
160                NSC_LOG_ERROR_WA("Failed to get bounds for array: ", tag);
161                return _T("UNKNOWN");
162        }
163        nCount = ( lUBound - lLBound + 1 );
164        std::wstring result;
165        for (int i=0; i < nCount; i++) {
166                std::wstring s = parse_item<T,U>(array[i]);
167                strEx::append_list(result, s);
168        }
169        SafeArrayUnaccessData(paArray);
170        return result;
171}
172
173WMIQuery::WMIResult get_value(std::wstring tag, CComVariant &vValue) {
174        WMIQuery::WMIResult value;
175        if (vValue.vt == VT_INT) {
176                value.setNumeric(vValue.intVal);
177        }else if (vValue.vt == VT_I2) {
178                value.setNumeric(vValue.iVal);
179        } else if (vValue.vt == VT_I4) {
180                value.setNumeric(vValue.lVal);
181        } else if (vValue.vt == VT_UI1) {
182                value.setNumeric(vValue.uintVal);
183        } else if (vValue.vt == VT_UINT) {
184                value.setNumeric(vValue.uintVal);
185        } else if (vValue.vt == VT_BSTR) {
186                value.setString(OLE2T(vValue.bstrVal));
187        } else if (vValue.vt == VT_NULL) {
188                value.setString(_T("NULL"));
189        } else if (vValue.vt == VT_BOOL) {
190                value.setBoth(vValue.iVal, vValue.iVal?_T("TRUE"):_T("FALSE"));
191        } else if (vValue.vt == (VT_ARRAY|VT_BSTR)) {
192                value.setString(array_to_string<BSTR,VT_BSTR>(tag, vValue));
193        } else if (vValue.vt == (VT_ARRAY|VT_BOOL)) {
194                value.setString(array_to_string<SHORT,VT_BOOL>(tag, vValue));
195        } else if (vValue.vt == (VT_ARRAY|VT_UI1)) {
196                value.setString(array_to_string<UINT,VT_UI1>(tag, vValue));
197        } else if (vValue.vt == (VT_ARRAY|VT_I4)) {
198                value.setString(array_to_string<LONG,VT_I4>(tag, vValue));
199        } else {
200                value.setString(_T("UNKNOWN"));
201                NSC_LOG_ERROR_W(tag + _T(" is not supported (type-id: ") + strEx::itos(vValue.vt) + _T(")"));
202        }
203        return value;
204}
205
206WMIQuery::result_type WMIQuery::get_classes(std::wstring ns, std::wstring superClass, std::wstring user, std::wstring password) {
207
208        result_type ret;
209        CComPtr<IWbemServices> service = create_service(ns, user, password);
210
211        CComBSTR strSuperClass = superClass.c_str();
212
213        CComPtr<IEnumWbemClassObject> enumerator;
214        HRESULT hr = service->CreateClassEnum(strSuperClass,WBEM_FLAG_DEEP| WBEM_FLAG_USE_AMENDED_QUALIFIERS|
215                WBEM_FLAG_RETURN_IMMEDIATELY|WBEM_FLAG_FORWARD_ONLY, NULL, &enumerator);
216        if (FAILED(hr))
217                throw WMIException("CreateClassEnum failed: " + ComError::getComError(ComError::getWMIError(hr)) + ")");
218
219        CComPtr<IWbemClassObject> row = NULL;
220        ULONG retcnt;
221        while (hr = enumerator->Next( WBEM_INFINITE, 1L, &row, &retcnt ) == WBEM_S_NO_ERROR) {
222                if (SUCCEEDED(hr) && (retcnt > 0)) {
223                        wmi_row returnRow;
224                        WMIResult value;
225                        CComBSTR classProp = _T("__CLASS");
226                        CComVariant v;
227                        hr = row->Get(classProp, 0, &v, 0, 0);
228
229                        if (SUCCEEDED(hr) && (v.vt == VT_BSTR)) {
230                                value.setString(std::wstring(OLE2T(v.bstrVal)));
231                                returnRow.addValue(_T("__CLASS"), value);
232                        } else {
233                                value.setString(_T("Unknown"));
234                                returnRow.addValue(_T("__CLASS"), value);
235                        }
236                        ret.push_back(returnRow);
237                }
238                row.Release();
239        }
240        return ret;
241}
242
243WMIQuery::result_type WMIQuery::get_instances(std::wstring ns, std::wstring superClass, std::wstring user, std::wstring password) {
244        result_type ret;
245        CComPtr<IWbemServices> service = create_service(ns, user, password);
246
247        CComBSTR strSuperClass = superClass.c_str();
248        CComPtr<IEnumWbemClassObject> enumerator;
249        HRESULT hr = service->CreateInstanceEnum(strSuperClass,WBEM_FLAG_SHALLOW|WBEM_FLAG_USE_AMENDED_QUALIFIERS|WBEM_FLAG_RETURN_IMMEDIATELY|WBEM_FLAG_FORWARD_ONLY,
250                NULL, &enumerator);
251        if (FAILED(hr))
252                throw WMIException("CreateInstanceEnum failed: " + ComError::getComError(ComError::getWMIError(hr)) + ")");
253
254
255        set_proxy_blanket(enumerator, user, password);
256        ULONG retcnt;
257        CComPtr< IWbemClassObject > row = NULL;
258        hr = enumerator->Next(WBEM_INFINITE, 1L, &row, &retcnt);
259        if (FAILED(hr))
260                throw WMIException("Enumeration of failed: " + ComError::getComError(ComError::getWMIError(hr)) + ")");
261
262        while (hr == WBEM_S_NO_ERROR) {
263                if (retcnt > 0) {
264                        wmi_row returnRow;
265                        WMIResult value;
266                        CComBSTR classProp = _T("Name");
267                        CComVariant v;
268                        hr = row->Get(classProp, 0, &v, 0, 0);
269
270                        if (SUCCEEDED(hr) && (v.vt == VT_BSTR)) {
271                                value.setString(std::wstring(OLE2T(v.bstrVal)));
272                                returnRow.addValue(_T("Name"), value);
273                        } else {
274                                value.setString(_T("Unknown"));
275                                returnRow.addValue(_T("Name"), value);
276                        }
277
278
279                        SAFEARRAY* pstrNames;
280                        hr = row->GetNames(NULL,WBEM_FLAG_ALWAYS|WBEM_FLAG_NONSYSTEM_ONLY,NULL,&pstrNames);
281                        if (FAILED(hr))
282                                throw WMIException("GetNames failed:", hr);
283
284                        long begin, end;
285                        CComSafeArray<BSTR> arr = pstrNames;
286                        begin = arr.GetLowerBound();
287                        end = arr.GetUpperBound();
288                        for (long index = begin; index <= end; index++ ) {
289                                try {
290                                        CComBSTR bColumn = arr.GetAt(index);
291                                        std::wstring column = bColumn.m_str;
292                                        CComVariant vValue;
293                                        hr = row->Get(bColumn, 0, &vValue, 0, 0);
294                                        if (FAILED(hr))
295                                                throw WMIException("Failed to get value for " + utf8::cvt<std::string>(column) + " in query: ", hr);
296                                        returnRow.addValue(column, get_value(column, vValue));
297                                } catch (const std::exception &e) {
298                                        throw WMIException("Failed to convert data: " + utf8::utf8_from_native(e.what()));
299                                }
300                        }
301
302                        ret.push_back(returnRow);
303                }
304                row.Release();
305                hr = enumerator->Next(WBEM_INFINITE, 1L, &row, &retcnt);
306                if (FAILED(hr))
307                        throw WMIException("Enumeration of failed: " + ComError::getComError(ComError::getWMIError(hr)) + ")");
308        }
309        return ret;
310}
311
312WMIQuery::result_type WMIQuery::execute(std::wstring ns, std::wstring query, std::wstring user, std::wstring password)
313{
314        result_type ret;
315        CComPtr<IWbemServices> service = create_service(ns, user, password);
316
317        CComBSTR strQuery(query.c_str());
318        BSTR strQL = _T("WQL");
319
320        CComPtr<IEnumWbemClassObject> enumerator;
321        HRESULT hr = service->ExecQuery( strQL, strQuery, WBEM_FLAG_FORWARD_ONLY, NULL, &enumerator );
322        if (FAILED(hr))
323                throw WMIException("ExecQuery of '" + utf8::cvt<std::string>(query) + "' failed: " + ComError::getComError(ComError::getWMIError(hr)));
324
325
326        set_proxy_blanket(enumerator, user, password);
327        ULONG retcnt;
328        CComPtr< IWbemClassObject > row = NULL;
329        hr = enumerator->Next(WBEM_INFINITE, 1L, &row, &retcnt);
330        if (FAILED(hr))
331                throw WMIException("Enumeration of '" + utf8::cvt<std::string>(query) + "' failed: " + ComError::getComError(ComError::getWMIError(hr)));
332        while (hr == WBEM_S_NO_ERROR) {
333                if (retcnt > 0) {
334                        SAFEARRAY* pstrNames;
335                        wmi_row returnRow;
336                        hr = row->GetNames(NULL,WBEM_FLAG_ALWAYS|WBEM_FLAG_NONSYSTEM_ONLY,NULL,&pstrNames);
337                        if (FAILED(hr))
338                                throw WMIException("GetNames failed:" + utf8::cvt<std::string>(query), hr);
339
340                        long begin, end;
341                        CComSafeArray<BSTR> arr = pstrNames;
342                        begin = arr.GetLowerBound();
343                        end = arr.GetUpperBound();
344                        for (long index = begin; index <= end; index++ ) {
345                                try {
346                                        CComBSTR bColumn = arr.GetAt(index);
347                                        std::wstring column = bColumn.m_str;
348                                        CComVariant vValue;
349                                        hr = row->Get(bColumn, 0, &vValue, 0, 0);
350                                        if (FAILED(hr))
351                                                throw WMIException("Failed to get value for " + utf8::cvt<std::string>(column) + " in query: " + utf8::cvt<std::string>(query), hr);
352                                        returnRow.addValue(column, get_value(column, vValue));
353                                } catch (const std::exception &e) {
354                                        throw WMIException("Failed to convert data: " + utf8::utf8_from_native(e.what()));
355                                }
356                        }
357                        ret.push_back(returnRow);
358                }
359                row.Release();
360                hr = enumerator->Next(WBEM_INFINITE, 1L, &row, &retcnt);
361                if (FAILED(hr))
362                        throw WMIException("Enumeration of '" + utf8::cvt<std::string>(query) + "' failed: " + ComError::getComError(ComError::getWMIError(hr)));
363        }
364        return ret;
365}
366
367std::string ComError::getComError(std::wstring inDesc /*= _T("")*/)
368{
369        return error::com::get();
370}
Note: See TracBrowser for help on using the repository browser.