Disk.cpp

#include "pch.h"
#include "Disk.h"
#include <Windows.h>
#include <random>

CDisk::CDisk(int index, const std::wstring& name, const std::wstring& type, ULONGLONG sizeInBytes)
    : m_nIndex(index), m_sName(name), m_sType(type)
{
    if (sizeInBytes > 0) {
        m_ullSizeInGB = sizeInBytes / 1024 / 1024 / 1024;
    }
}

bool CDisk::IsRemovable() const
{
    // Simple case-insensitive comparison
    return _wcsicmp(m_sType.c_str(), L"Removable Media") == 0;
}

void CDisk::AssignEFIPartitionLetter()
{
    std::wstring availableLetters = L"ABCDEFGHIJKLMNOPQRSTUVWXYZ";
    std::wstring usedLetters;

    DWORD drives = GetLogicalDrives();
    for (char i = 0; i < 26; ++i) {
        if ((drives >> i) & 1) {
            usedLetters += (wchar_t)(L'A' + i);
        }
    }

    // Remove used letters from available letters
    for (wchar_t used : usedLetters) {
        size_t pos = availableLetters.find(used);
        if (pos != std::wstring::npos) {
            availableLetters.erase(pos, 1);
        }
    }

    if (!availableLetters.empty()) {
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> distrib(0, availableLetters.length() - 1);
        m_sEFIPartitionLetter = availableLetters[distrib(gen)];
    }
    else {
        // Fallback or throw an exception if no letters are available
        m_sEFIPartitionLetter.clear();
    }
}