master
c 144 lines 5.28 KB
Raw
1 // SPDX-License-Identifier: GPL-3.0-or-later
2
3 #include "windows-wmi-GetDiskDriveInfo.h"
4
5 #if defined(OS_WINDOWS)
6
7 size_t GetDiskDriveInfo(DiskDriveInfoWMI *diskInfoArray, size_t array_size) {
8 if (InitializeWMI() != S_OK) return 0;
9
10 HRESULT hr;
11 IEnumWbemClassObject* pEnumerator = NULL;
12
13 // Execute the query, including new properties
14 BSTR query = SysAllocString(L"SELECT DeviceID, Model, Caption, Name, Partitions, Size, Status, Availability, Index, Manufacturer, InstallDate, MediaType, NeedsCleaning FROM WIN32_DiskDrive");
15 BSTR wql = SysAllocString(L"WQL");
16 hr = nd_wmi.pSvc->lpVtbl->ExecQuery(
17 nd_wmi.pSvc,
18 wql,
19 query,
20 WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
21 NULL,
22 &pEnumerator
23 );
24 SysFreeString(query);
25 SysFreeString(wql);
26
27 if (FAILED(hr)) {
28 nd_log(NDLS_COLLECTORS, NDLP_ERR, "GetDiskDriveInfo() WMI query failed. Error code = 0x%X", hr);
29 return 0;
30 }
31
32 // Iterate through the results
33 IWbemClassObject *pclsObj = NULL;
34 ULONG uReturn = 0;
35 size_t index = 0;
36 while (pEnumerator && index < array_size) {
37 hr = pEnumerator->lpVtbl->Next(pEnumerator, WBEM_INFINITE, 1, &pclsObj, &uReturn);
38 if (0 == uReturn) break;
39
40 VARIANT vtProp;
41
42 // Extract DeviceID
43 hr = pclsObj->lpVtbl->Get(pclsObj, L"DeviceID", 0, &vtProp, 0, 0);
44 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
45 wcstombs(diskInfoArray[index].DeviceID, vtProp.bstrVal, sizeof(diskInfoArray[index].DeviceID));
46 }
47 VariantClear(&vtProp);
48
49 // Extract Model
50 hr = pclsObj->lpVtbl->Get(pclsObj, L"Model", 0, &vtProp, 0, 0);
51 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
52 wcstombs(diskInfoArray[index].Model, vtProp.bstrVal, sizeof(diskInfoArray[index].Model));
53 }
54 VariantClear(&vtProp);
55
56 // Extract Caption
57 hr = pclsObj->lpVtbl->Get(pclsObj, L"Caption", 0, &vtProp, 0, 0);
58 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
59 wcstombs(diskInfoArray[index].Caption, vtProp.bstrVal, sizeof(diskInfoArray[index].Caption));
60 }
61 VariantClear(&vtProp);
62
63 // Extract Name
64 hr = pclsObj->lpVtbl->Get(pclsObj, L"Name", 0, &vtProp, 0, 0);
65 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
66 wcstombs(diskInfoArray[index].Name, vtProp.bstrVal, sizeof(diskInfoArray[index].Name));
67 }
68 VariantClear(&vtProp);
69
70 // Extract Partitions
71 hr = pclsObj->lpVtbl->Get(pclsObj, L"Partitions", 0, &vtProp, 0, 0);
72 if (SUCCEEDED(hr) && (vtProp.vt == VT_I4 || vtProp.vt == VT_UI4)) {
73 diskInfoArray[index].Partitions = vtProp.intVal;
74 }
75 VariantClear(&vtProp);
76
77 // Extract Size (convert BSTR to uint64)
78 hr = pclsObj->lpVtbl->Get(pclsObj, L"Size", 0, &vtProp, 0, 0);
79 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
80 char sizeStr[64];
81 wcstombs(sizeStr, vtProp.bstrVal, sizeof(sizeStr));
82 diskInfoArray[index].Size = strtoull(sizeStr, NULL, 10);
83 }
84 VariantClear(&vtProp);
85
86 // Extract Status
87 hr = pclsObj->lpVtbl->Get(pclsObj, L"Status", 0, &vtProp, 0, 0);
88 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
89 wcstombs(diskInfoArray[index].Status, vtProp.bstrVal, sizeof(diskInfoArray[index].Status));
90 }
91 VariantClear(&vtProp);
92
93 // Extract Availability
94 hr = pclsObj->lpVtbl->Get(pclsObj, L"Availability", 0, &vtProp, 0, 0);
95 if (SUCCEEDED(hr) && (vtProp.vt == VT_I4 || vtProp.vt == VT_UI4)) {
96 diskInfoArray[index].Availability = vtProp.intVal;
97 }
98 VariantClear(&vtProp);
99
100 // Extract Index
101 hr = pclsObj->lpVtbl->Get(pclsObj, L"Index", 0, &vtProp, 0, 0);
102 if (SUCCEEDED(hr) && (vtProp.vt == VT_I4 || vtProp.vt == VT_UI4)) {
103 diskInfoArray[index].Index = vtProp.intVal;
104 }
105 VariantClear(&vtProp);
106
107 // Extract Manufacturer
108 hr = pclsObj->lpVtbl->Get(pclsObj, L"Manufacturer", 0, &vtProp, 0, 0);
109 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
110 wcstombs(diskInfoArray[index].Manufacturer, vtProp.bstrVal, sizeof(diskInfoArray[index].Manufacturer));
111 }
112 VariantClear(&vtProp);
113
114 // Extract InstallDate
115 hr = pclsObj->lpVtbl->Get(pclsObj, L"InstallDate", 0, &vtProp, 0, 0);
116 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
117 wcstombs(diskInfoArray[index].InstallDate, vtProp.bstrVal, sizeof(diskInfoArray[index].InstallDate));
118 }
119 VariantClear(&vtProp);
120
121 // Extract MediaType
122 hr = pclsObj->lpVtbl->Get(pclsObj, L"MediaType", 0, &vtProp, 0, 0);
123 if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
124 wcstombs(diskInfoArray[index].MediaType, vtProp.bstrVal, sizeof(diskInfoArray[index].MediaType));
125 }
126 VariantClear(&vtProp);
127
128 // Extract NeedsCleaning
129 hr = pclsObj->lpVtbl->Get(pclsObj, L"NeedsCleaning", 0, &vtProp, 0, 0);
130 if (SUCCEEDED(hr) && (vtProp.vt == VT_BOOL)) {
131 diskInfoArray[index].NeedsCleaning = vtProp.boolVal;
132 }
133 VariantClear(&vtProp);
134
135 pclsObj->lpVtbl->Release(pclsObj);
136 index++;
137 }
138
139 pEnumerator->lpVtbl->Release(pEnumerator);
140
141 return index;
142 }
143
144 #endif