WmiHelper.cpp

#include "pch.h"
#include "WmiHelper.h"
#include <Wbemidl.h>
#include <comdef.h>
#include <stdexcept>
#include <winioctl.h>
#pragma comment(lib, "wbemuuid.lib")

// {C12A7328-F81F-11D2-BA4B-00A0C93EC93B}: type GUID of EFI System Partition
const GUID PARTITION_SYSTEM_GUID =
{ 0xC12A7328, 0xF81F, 0x11D2, { 0xBA, 0x4B, 0x00, 0xA0, 0xC9, 0x3E, 0xC9, 0x3B } };

bool WmiHelper::GetDisksList(std::vector<CDisk>& disks)
{
    HRESULT hres;

    IWbemLocator* pLoc = NULL;
    hres = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID*)&pLoc);
    if (FAILED(hres)) {
        CoUninitialize();
        throw std::runtime_error("Failed to create IWbemLocator object.");
    }

    IWbemServices* pSvc = NULL;
    hres = pLoc->ConnectServer(_bstr_t(L"ROOT\\CIMV2"), NULL, NULL, 0, NULL, 0, 0, &pSvc);
    if (FAILED(hres)) {
        pLoc->Release();
        CoUninitialize();
        throw std::runtime_error("Could not connect to WMI server.");
    }

    hres = CoSetProxyBlanket(
        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(hres)) {
        pSvc->Release();
        pLoc->Release();
        CoUninitialize();
        throw std::runtime_error("Could not set proxy blanket.");
    }

    IEnumWbemClassObject* pEnumerator = NULL;
    hres = pSvc->ExecQuery(
        bstr_t("WQL"),
        bstr_t("SELECT * FROM Win32_DiskDrive"),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        NULL,
        &pEnumerator);
    if (FAILED(hres)) {
        pSvc->Release();
        pLoc->Release();
        CoUninitialize();
        throw std::runtime_error("WMI query for disk drives failed.");
    }

    IWbemClassObject* pclsObj = NULL;
    ULONG uReturn = 0;

    while (pEnumerator)
    {
        HRESULT hr = pEnumerator->Next(WBEM_INFINITE, 1, &pclsObj, &uReturn);
        if (0 == uReturn) break;

        VARIANT vtProp;
        int index = -1;
        std::wstring model = L"Unknown Model";
        std::wstring mediaType = L"Unknown Media Type";
        ULONGLONG size = 0;

        hr = pclsObj->Get(L"Index", 0, &vtProp, 0, 0);
        if (SUCCEEDED(hr) && vtProp.vt == VT_I4) index = vtProp.intVal;
        VariantClear(&vtProp);

        hr = pclsObj->Get(L"Model", 0, &vtProp, 0, 0);
        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) model = vtProp.bstrVal;
        VariantClear(&vtProp);

        hr = pclsObj->Get(L"MediaType", 0, &vtProp, 0, 0);
        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) mediaType = vtProp.bstrVal;
        VariantClear(&vtProp);

        hr = pclsObj->Get(L"Size", 0, &vtProp, 0, 0);
        if (SUCCEEDED(hr) && vtProp.vt == VT_BSTR) {
            size = _wtoi64(vtProp.bstrVal);
        }
        else if (SUCCEEDED(hr) && vtProp.vt == VT_UI8) { // Sometimes it's a numeric type
            size = vtProp.ullVal;
        }
        VariantClear(&vtProp);

        // Only add the disk to the list if it contains an EFI partition
        if (index != -1) {
            int efiPartitionIndex = FindEfiPartitionIndex(index);
            if (efiPartitionIndex > 0) {
                disks.emplace_back(index, model, mediaType, size);
                disks.back().m_nEfiPartitionIndex = efiPartitionIndex; // Store EFI partition index
            }
        }

        pclsObj->Release();
    }

    pSvc->Release();
    pLoc->Release();
    pEnumerator->Release();

    return true;
}

int WmiHelper::FindEfiPartitionIndex(int diskIndex) {
    std::wstring path = L"\\\\.\\PhysicalDrive" + std::to_wstring(diskIndex);
    HANDLE hDevice = CreateFileW(path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, 0, NULL);

    if (hDevice == INVALID_HANDLE_VALUE) {
        return false;
    }

    // To accommodate the partition layout information, allocate a buffer large enough (supports up to 128 partitions)
    std::vector<BYTE> buffer(sizeof(DRIVE_LAYOUT_INFORMATION_EX) + 127 * sizeof(PARTITION_INFORMATION_EX));
    DRIVE_LAYOUT_INFORMATION_EX* pLayout = reinterpret_cast<DRIVE_LAYOUT_INFORMATION_EX*>(buffer.data());
    DWORD bytesReturned = 0;

    bool result = DeviceIoControl(
        hDevice,
        IOCTL_DISK_GET_DRIVE_LAYOUT_EX,
        NULL, 0,
        pLayout, (DWORD)buffer.size(),
        &bytesReturned,
        NULL
    );

    int efiIndex = -1;
    if (result && pLayout->PartitionStyle == PARTITION_STYLE_GPT) {
        for (DWORD i = 0; i < pLayout->PartitionCount; i++) {
            if (IsEqualGUID(pLayout->PartitionEntry[i].Gpt.PartitionType, PARTITION_SYSTEM_GUID)) {
                efiIndex = i + 1;
                break;
            }
        }
    }

    CloseHandle(hDevice);
    return efiIndex;
}