Merging upstream version 2.0.3+dfsg (Closes: #923993, #1042533, #1045145).

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

4 files changed, 307 insertions, 0 deletions

diff --git a/src/libnetdata/os/windows-wmi/windows-wmi-GetDiskDriveInfo.c b/src/libnetdata/os/windows-wmi/windows-wmi-GetDiskDriveInfo.c
new file mode 100644
index 000000000..283c6f09e
--- /dev/null
+++ b/src/libnetdata/os/windows-wmi/windows-wmi-GetDiskDriveInfo.c
@@ -0,0 +1,144 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "windows-wmi-GetDiskDriveInfo.h"
+
+#if defined(OS_WINDOWS)
+
+size_t GetDiskDriveInfo(DiskDriveInfoWMI *diskInfoArray, size_t array_size) {
+    if (InitializeWMI() != S_OK) return 0;
+
+    HRESULT hr;
+    IEnumWbemClassObject* pEnumerator = NULL;
+
+    // Execute the query, including new properties
+    BSTR query = SysAllocString(L"SELECT DeviceID, Model, Caption, Name, Partitions, Size, Status, Availability, Index, Manufacturer, InstallDate, MediaType, NeedsCleaning FROM WIN32_DiskDrive");
+    BSTR wql = SysAllocString(L"WQL");
+    hr = nd_wmi.pSvc->lpVtbl->ExecQuery(
+            nd_wmi.pSvc,
+            wql,
+            query,
+            WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
+            NULL,
+            &pEnumerator
+    );
+    SysFreeString(query);
+    SysFreeString(wql);
+
+    if (FAILED(hr)) {
+        nd_log(NDLS_COLLECTORS, NDLP_ERR, "GetDiskDriveInfo() WMI query failed. Error code = 0x%X", hr);
+        return 0;
+    }
+
+    // Iterate through the results
+    IWbemClassObject *pclsObj = NULL;
+    ULONG uReturn = 0;
+    size_t index = 0;
+    while (pEnumerator && index < array_size) {
+        hr = pEnumerator->lpVtbl->Next(pEnumerator, WBEM_INFINITE, 1, &pclsObj, &uReturn);
+        if (0 == uReturn) break;
+
+        VARIANT vtProp;
+
+        // Extract DeviceID
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"DeviceID", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].DeviceID, vtProp.bstrVal, sizeof(diskInfoArray[index].DeviceID));
+        }
+        VariantClear(&vtProp);
+
+        // Extract Model
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Model", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].Model, vtProp.bstrVal, sizeof(diskInfoArray[index].Model));
+        }
+        VariantClear(&vtProp);
+
+        // Extract Caption
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Caption", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].Caption, vtProp.bstrVal, sizeof(diskInfoArray[index].Caption));
+        }
+        VariantClear(&vtProp);
+
+        // Extract Name
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Name", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].Name, vtProp.bstrVal, sizeof(diskInfoArray[index].Name));
+        }
+        VariantClear(&vtProp);
+
+        // Extract Partitions
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Partitions", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && (vtProp.vt == VT_I4 || vtProp.vt == VT_UI4)) {
+            diskInfoArray[index].Partitions = vtProp.intVal;
+        }
+        VariantClear(&vtProp);
+
+        // Extract Size (convert BSTR to uint64)
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Size", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            char sizeStr[64];
+            wcstombs(sizeStr, vtProp.bstrVal, sizeof(sizeStr));
+            diskInfoArray[index].Size = strtoull(sizeStr, NULL, 10);
+        }
+        VariantClear(&vtProp);
+
+        // Extract Status
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Status", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].Status, vtProp.bstrVal, sizeof(diskInfoArray[index].Status));
+        }
+        VariantClear(&vtProp);
+
+        // Extract Availability
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Availability", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && (vtProp.vt == VT_I4 || vtProp.vt == VT_UI4)) {
+            diskInfoArray[index].Availability = vtProp.intVal;
+        }
+        VariantClear(&vtProp);
+
+        // Extract Index
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Index", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && (vtProp.vt == VT_I4 || vtProp.vt == VT_UI4)) {
+            diskInfoArray[index].Index = vtProp.intVal;
+        }
+        VariantClear(&vtProp);
+
+        // Extract Manufacturer
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"Manufacturer", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].Manufacturer, vtProp.bstrVal, sizeof(diskInfoArray[index].Manufacturer));
+        }
+        VariantClear(&vtProp);
+
+        // Extract InstallDate
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"InstallDate", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].InstallDate, vtProp.bstrVal, sizeof(diskInfoArray[index].InstallDate));
+        }
+        VariantClear(&vtProp);
+
+        // Extract MediaType
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"MediaType", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
+            wcstombs(diskInfoArray[index].MediaType, vtProp.bstrVal, sizeof(diskInfoArray[index].MediaType));
+        }
+        VariantClear(&vtProp);
+
+        // Extract NeedsCleaning
+        hr = pclsObj->lpVtbl->Get(pclsObj, L"NeedsCleaning", 0, &vtProp, 0, 0);
+        if (SUCCEEDED(hr) && (vtProp.vt == VT_BOOL)) {
+            diskInfoArray[index].NeedsCleaning = vtProp.boolVal;
+        }
+        VariantClear(&vtProp);
+
+        pclsObj->lpVtbl->Release(pclsObj);
+        index++;
+    }
+
+    pEnumerator->lpVtbl->Release(pEnumerator);
+
+    return index;
+}
+
+#endif
\ No newline at end of file
diff --git a/src/libnetdata/os/windows-wmi/windows-wmi-GetDiskDriveInfo.h b/src/libnetdata/os/windows-wmi/windows-wmi-GetDiskDriveInfo.h
new file mode 100644
index 000000000..cc9b46067
--- /dev/null
+++ b/src/libnetdata/os/windows-wmi/windows-wmi-GetDiskDriveInfo.h
@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef NETDATA_WINDOWS_WMI_GETDISKDRIVEINFO_H
+#define NETDATA_WINDOWS_WMI_GETDISKDRIVEINFO_H
+
+#include "windows-wmi.h"
+
+#if defined(OS_WINDOWS)
+
+typedef struct {
+    char DeviceID[256];
+    char Model[256];
+    char Caption[256];
+    char Name[256];
+    int Partitions;
+    unsigned long long Size;
+    char Status[64];
+    int Availability;
+    int Index;
+    char Manufacturer[256];
+    char InstallDate[64];
+    char MediaType[128];
+    bool NeedsCleaning;
+} DiskDriveInfoWMI;
+
+size_t GetDiskDriveInfo(DiskDriveInfoWMI *diskInfoArray, size_t array_size);
+
+#endif
+
+#endif //NETDATA_WINDOWS_WMI_GETDISKDRIVEINFO_H
diff --git a/src/libnetdata/os/windows-wmi/windows-wmi.c b/src/libnetdata/os/windows-wmi/windows-wmi.c
new file mode 100644
index 000000000..02d3faa7c
--- /dev/null
+++ b/src/libnetdata/os/windows-wmi/windows-wmi.c
@@ -0,0 +1,110 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "windows-wmi.h"
+
+#if defined(OS_WINDOWS)
+
+__thread ND_WMI nd_wmi = { 0 };
+
+HRESULT InitializeWMI(void) {
+    if(nd_wmi.pLoc && nd_wmi.pSvc) return S_OK;
+    CleanupWMI();
+
+    IWbemLocator **pLoc = &nd_wmi.pLoc;
+    IWbemServices **pSvc = &nd_wmi.pSvc;
+
+    HRESULT hr;
+
+    // Initialize COM
+    hr = CoInitializeEx(0, COINIT_MULTITHREADED);
+    if (FAILED(hr)) {
+        nd_log(NDLS_COLLECTORS, NDLP_ERR, "Failed to initialize COM library. Error code = 0x%X", hr);
+        CleanupWMI();
+        return hr;
+    }
+
+    // Set COM security levels
+    hr = CoInitializeSecurity(
+            NULL,
+            -1,
+            NULL,
+            NULL,
+            RPC_C_AUTHN_LEVEL_DEFAULT,
+            RPC_C_IMP_LEVEL_IMPERSONATE,
+            NULL,
+            EOAC_NONE,
+            NULL
+    );
+    if (FAILED(hr)) {
+        nd_log(NDLS_COLLECTORS, NDLP_ERR, "Failed to initialize security. Error code = 0x%X", hr);
+        CleanupWMI();
+        return hr;
+    }
+
+    // Obtain the initial locator to WMI
+    hr = CoCreateInstance(
+            &CLSID_WbemLocator, 0,
+            CLSCTX_INPROC_SERVER,
+            &IID_IWbemLocator, (LPVOID *)pLoc
+    );
+    if (FAILED(hr)) {
+        nd_log(NDLS_COLLECTORS, NDLP_ERR, "Failed to create IWbemLocator object. Error code = 0x%X", hr);
+        CleanupWMI();
+        return hr;
+    }
+
+    // Connect to WMI
+    BSTR namespacePath = SysAllocString(L"ROOT\\CIMV2");
+    hr = (*pLoc)->lpVtbl->ConnectServer(
+            *pLoc,
+            namespacePath,
+            NULL,
+            NULL,
+            0,
+            0,
+            0,
+            0,
+            pSvc
+    );
+    SysFreeString(namespacePath);
+
+    if (FAILED(hr)) {
+        nd_log(NDLS_COLLECTORS, NDLP_ERR, "Could not connect to WMI server. Error code = 0x%X", hr);
+        CleanupWMI();
+        return hr;
+    }
+
+    // Set security levels on the proxy
+    hr = CoSetProxyBlanket(
+            (IUnknown *)*pSvc,
+            RPC_C_AUTHN_WINNT,
+            RPC_C_AUTHZ_NONE,
+            NULL,
+            RPC_C_AUTHN_LEVEL_CALL,
+            RPC_C_IMP_LEVEL_IMPERSONATE,
+            NULL,
+            EOAC_NONE
+    );
+    if (FAILED(hr)) {
+        nd_log(NDLS_COLLECTORS, NDLP_ERR, "Could not set proxy blanket. Error code = 0x%X", hr);
+        CleanupWMI();
+        return hr;
+    }
+
+    return S_OK;
+}
+
+void CleanupWMI(void) {
+    if(nd_wmi.pLoc)
+        nd_wmi.pLoc->lpVtbl->Release(nd_wmi.pLoc);
+
+    if (nd_wmi.pSvc)
+        nd_wmi.pSvc->lpVtbl->Release(nd_wmi.pSvc);
+
+    nd_wmi.pLoc = NULL;
+    nd_wmi.pSvc = NULL;
+
+    CoUninitialize();
+}
+
+#endif
\ No newline at end of file
diff --git a/src/libnetdata/os/windows-wmi/windows-wmi.h b/src/libnetdata/os/windows-wmi/windows-wmi.h
new file mode 100644
index 000000000..69d7244aa
--- /dev/null
+++ b/src/libnetdata/os/windows-wmi/windows-wmi.h
@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef NETDATA_WINDOWS_WMI_H
+#define NETDATA_WINDOWS_WMI_H
+
+#include "../../libnetdata.h"
+
+#if defined(OS_WINDOWS)
+typedef struct {
+    IWbemLocator *pLoc;
+    IWbemServices *pSvc;
+} ND_WMI;
+
+extern __thread ND_WMI nd_wmi;
+
+HRESULT InitializeWMI(void);
+void CleanupWMI(void);
+
+#include "windows-wmi-GetDiskDriveInfo.h"
+
+#endif
+
+#endif //NETDATA_WINDOWS_WMI_H