EFIMounterDlg.cpp

// EFIMounterDlg.cpp: archivo de implementación
//

#include "pch.h"
#include "EFIMounter.h"
#include "EFIMounterDlg.h"
#include "afxdialogex.h"
#include "WmiHelper.h"
#include "Diskpart.h"
#include <algorithm>

#ifdef _DEBUG
#define new DEBUG_NEW
#endif


// Cuadro de diálogo de CEFIMounterDlg

// Worker thread functions
void LoadDisksThread(HWND hWnd) {
    // Every thread that uses COM must initialize its own COM objects!
    if (SUCCEEDED(CoInitializeEx(NULL, COINIT_MULTITHREADED)))
    {
        try {
            auto disks = new std::vector<CDisk>();
            WmiHelper::GetDisksList(*disks);
            std::sort(disks->begin(), disks->end(), [](const CDisk& a, const CDisk& b) {
                return a.m_nIndex > b.m_nIndex;
                });
            ::PostMessage(hWnd, WM_DISK_LOAD_COMPLETE, TRUE, (LPARAM)disks);
        }
        catch (const std::exception& e) {
            CString* errorMsg = new CString(e.what());
            ::PostMessage(hWnd, WM_DISK_LOAD_COMPLETE, FALSE, (LPARAM)errorMsg);
        }
        // Before the thread ends, clean up its own COM environment
        CoUninitialize();
    }
    else
    {
        CString* errorMsg = new CString(L"Failed to initialize COM library in worker thread.");
        ::PostMessage(hWnd, WM_DISK_LOAD_COMPLETE, FALSE, (LPARAM)errorMsg);
    }
}

void MountThread(HWND hWnd, int diskIndex, std::vector<CDisk>* disks, bool updateGuid) {
	try {
		(*disks)[diskIndex].AssignEFIPartitionLetter();
		Diskpart::AssignLetterToEFIPartition((*disks)[diskIndex], updateGuid);
		::PostMessage(hWnd, WM_MOUNT_COMPLETE, TRUE, diskIndex);
	}
	catch (const std::exception& e) {
		CString* errorMsg = new CString(e.what());
		::PostMessage(hWnd, WM_MOUNT_COMPLETE, FALSE, (LPARAM)errorMsg);
	}
}

void UnmountThread(HWND hWnd, int diskIndex, std::vector<CDisk>* disks, bool updateGuid) {
	try {
		Diskpart::RemoveEFIPartitionLetter((*disks)[diskIndex], updateGuid);
		(*disks)[diskIndex].m_sEFIPartitionLetter.clear();
		::PostMessage(hWnd, WM_UNMOUNT_COMPLETE, TRUE, diskIndex);
	}
	catch (const std::exception& e) {
		CString* errorMsg = new CString(e.what());
		::PostMessage(hWnd, WM_UNMOUNT_COMPLETE, FALSE, (LPARAM)errorMsg);
	}
}

CEFIMounterDlg::CEFIMounterDlg(CWnd* pParent /*= nullptr*/)
	: CDialogEx(IDD_EFIMOUNTER_DIALOG, pParent)
{
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CEFIMounterDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialogEx::DoDataExchange(pDX);
	DDX_Control(pDX, IDC_DRIVES_COMBO, m_comboDrives);
	DDX_Control(pDX, IDC_REFRESH_BTN, m_btnRefresh);
	DDX_Control(pDX, IDC_EXPLORE_BTN, m_btnExplore);
	DDX_Control(pDX, IDC_MOUNT_BTN, m_btnMount);
	DDX_Control(pDX, IDC_UNMOUNT_BTN, m_btnUnmount);
	DDX_Control(pDX, IDC_UPDATE_GUID_CHECK, m_checkUpdateGuid);
	DDX_Control(pDX, IDC_PROGRESS, m_progress);
	DDX_Control(pDX, IDC_STATUS_BAR, m_statusBar);
}

BEGIN_MESSAGE_MAP(CEFIMounterDlg, CDialogEx)
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	ON_BN_CLICKED(IDC_REFRESH_BTN, &CEFIMounterDlg::OnBnClickedRefreshBtn)
	ON_CBN_SELCHANGE(IDC_DRIVES_COMBO, &CEFIMounterDlg::OnCbnSelchangeDrivesCombo)
	ON_BN_CLICKED(IDC_EXPLORE_BTN, &CEFIMounterDlg::OnBnClickedExploreBtn)
	ON_BN_CLICKED(IDC_MOUNT_BTN, &CEFIMounterDlg::OnBnClickedMountBtn)
	ON_BN_CLICKED(IDC_UNMOUNT_BTN, &CEFIMounterDlg::OnBnClickedUnmountBtn)
	ON_WM_CLOSE()
	ON_MESSAGE(WM_DISK_LOAD_COMPLETE, &CEFIMounterDlg::OnDiskLoadComplete)
	ON_MESSAGE(WM_MOUNT_COMPLETE, &CEFIMounterDlg::OnMountComplete)
	ON_MESSAGE(WM_UNMOUNT_COMPLETE, &CEFIMounterDlg::OnUnmountComplete)
END_MESSAGE_MAP()


// Controladores de mensajes de CEFIMounterDlg

BOOL CEFIMounterDlg::OnInitDialog()
{
	CDialogEx::OnInitDialog();

	SetIcon(m_hIcon, TRUE);
	SetIcon(m_hIcon, FALSE);

	// Set symbol font for refresh button
	m_symbolFont.CreatePointFont(120, L"Segoe MDL2 Assets");
	m_btnRefresh.SetFont(&m_symbolFont);
    m_btnRefresh.SetWindowTextW(L"\uE149"); // refresh logo

    // Aunque la barra de estado se crea en el recurso, necesitamos definir sus paneles.
    // Crearemos un único panel expandible.
    UINT indicators[] = { ID_SEPARATOR };
    m_statusBar.SetIndicators(indicators, 1);

    // Hacemos que el primer panel (índice 0) se expanda para ocupar el espacio disponible.
    m_statusBar.SetPaneInfo(0, ID_SEPARATOR, SBPS_STRETCH, 0);

    CoInitializeEx(NULL, COINIT_APARTMENTTHREADED);

    // Set up COM security for the entire process. This only needs to be called once!
    HRESULT hres = CoInitializeSecurity(
        NULL,
        -1,                          // COM negotiates service
        NULL,                        // Authentication services
        NULL,                        // Reserved
        RPC_C_AUTHN_LEVEL_DEFAULT,   // Default authentication
        RPC_C_IMP_LEVEL_IMPERSONATE, // Default Impersonation
        NULL,                        // Authentication info
        EOAC_NONE,                   // Additional capabilities
        NULL                         // Reserved
    );
    if (FAILED(hres)) {
        MessageBox(L"Failed to initialize COM security.", L"Fatal Error", MB_OK | MB_ICONERROR);
        OnCancel(); // Exit directly on fail
        return FALSE;
    }

	LoadDisksAsync();

	return TRUE;
}

void CEFIMounterDlg::OnPaint()
{
	if (IsIconic())
	{
		CPaintDC dc(this); // Contexto de dispositivo para dibujo

		SendMessage(WM_ICONERASEBKGND, reinterpret_cast<WPARAM>(dc.GetSafeHdc()), 0);

		// Centrar icono en el rect치ngulo de cliente
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// Dibujar el icono
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialogEx::OnPaint();
	}
}

// El sistema llama a esta función para obtener el cursor que se muestra mientras el usuario arrastra
//  la ventana minimizada.
HCURSOR CEFIMounterDlg::OnQueryDragIcon()
{
	return static_cast<HCURSOR>(m_hIcon);
}

void CEFIMounterDlg::EnableControls(BOOL bEnable)
{
    m_comboDrives.EnableWindow(bEnable);
    m_btnRefresh.EnableWindow(bEnable);
    m_btnMount.EnableWindow(bEnable);
    m_btnUnmount.EnableWindow(bEnable);
    m_btnExplore.EnableWindow(bEnable);
    m_checkUpdateGuid.EnableWindow(bEnable);
}

void CEFIMounterDlg::SetStatus(const CString& sMessage, bool bProgress)
{
    m_statusBar.SetPaneText(0, sMessage);

    if (bProgress)
    {
        m_progress.ShowWindow(SW_SHOW);
        m_progress.SetMarquee(TRUE, 30);
    }
    else
    {
        m_progress.SetMarquee(FALSE, 0);
        m_progress.ShowWindow(SW_HIDE);
    }
    m_progress.UpdateWindow();

    if (bProgress)
    {
        EnableControls(FALSE);
    }
}

void CEFIMounterDlg::LoadDisksAsync()
{
    SetStatus(L"Loading drives...", true);
    m_comboDrives.ResetContent();
    m_disks.clear();
    m_nCurrentDiskIndex = -1;
    m_nPreviousDiskIndex = -1;

    std::thread(LoadDisksThread, m_hWnd).detach();
}

void CEFIMounterDlg::OnBnClickedRefreshBtn()
{
    if (m_nCurrentDiskIndex != -1 && !m_disks[m_nCurrentDiskIndex].m_sEFIPartitionLetter.empty()) {
        UnmountAsync(m_nCurrentDiskIndex);
    }
    LoadDisksAsync();
}

void CEFIMounterDlg::OnCbnSelchangeDrivesCombo()
{
    m_nPreviousDiskIndex = m_nCurrentDiskIndex;
    m_nCurrentDiskIndex = m_comboDrives.GetCurSel();

    if (m_nPreviousDiskIndex != -1 && m_nPreviousDiskIndex != m_nCurrentDiskIndex &&
        !m_disks[m_nPreviousDiskIndex].m_sEFIPartitionLetter.empty())
    {
        UnmountAsync(m_nPreviousDiskIndex);
    }
    UpdateUIStateForSelection();
}

void CEFIMounterDlg::UpdateUIStateForSelection()
{
    if (m_nCurrentDiskIndex < 0 || m_nCurrentDiskIndex >= m_disks.size()) {
        m_btnMount.EnableWindow(FALSE);
        m_btnUnmount.EnableWindow(FALSE);
        m_btnExplore.EnableWindow(FALSE);
        m_checkUpdateGuid.EnableWindow(FALSE);
        SetStatus(L"No drive selected", false);
        return;
    }

    const auto& selectedDisk = m_disks[m_nCurrentDiskIndex];
    bool isMounted = !selectedDisk.m_sEFIPartitionLetter.empty();

    m_btnMount.EnableWindow(!isMounted);
    m_btnUnmount.EnableWindow(isMounted);
    m_btnExplore.EnableWindow(TRUE);

    bool guidCheckEnabled = !isMounted && selectedDisk.IsRemovable();
    m_checkUpdateGuid.EnableWindow(guidCheckEnabled);
    m_checkUpdateGuid.SetCheck(guidCheckEnabled ? BST_CHECKED : BST_UNCHECKED);

    if (isMounted) {
        CString status;
        status.Format(L"EFI partition mounted on %s:", selectedDisk.m_sEFIPartitionLetter.c_str());
        SetStatus(status, false);
    }
    else {
        SetStatus(L"EFI partition is not mounted", false);
    }
}

void CEFIMounterDlg::MountAsync(int diskIndex)
{
    if (diskIndex < 0 || diskIndex >= m_disks.size()) return;
    if (!m_disks[diskIndex].m_sEFIPartitionLetter.empty()) return;

    SetStatus(L"Mounting the EFI partition...", true);
    bool updateGuid = m_checkUpdateGuid.GetCheck() == BST_CHECKED;
    std::thread(MountThread, m_hWnd, diskIndex, &m_disks, updateGuid).detach();
}

void CEFIMounterDlg::UnmountAsync(int diskIndex)
{
    if (diskIndex < 0 || diskIndex >= m_disks.size()) return;
    if (m_disks[diskIndex].m_sEFIPartitionLetter.empty()) return;

    SetStatus(L"Unmounting the EFI partition...", true);
    bool updateGuid = m_checkUpdateGuid.GetCheck() == BST_CHECKED;
    std::thread(UnmountThread, m_hWnd, diskIndex, &m_disks, updateGuid).detach();
}

void CEFIMounterDlg::OnBnClickedMountBtn()
{
    MountAsync(m_nCurrentDiskIndex);
}

void CEFIMounterDlg::OnBnClickedUnmountBtn()
{
    UnmountAsync(m_nCurrentDiskIndex);
}

void CEFIMounterDlg::OnBnClickedExploreBtn()
{
    if (m_nCurrentDiskIndex < 0 || m_nCurrentDiskIndex >= m_disks.size()) return;

    if (m_disks[m_nCurrentDiskIndex].m_sEFIPartitionLetter.empty()) {
        // Set the flag and start mounting if the partition is not mounted
        m_bExploreAfterMount = true;
        MountAsync(m_nCurrentDiskIndex);
    }
    else {
        // Otherwise directly pops up the file selection dialog box
        CString initialDir = m_disks[m_nCurrentDiskIndex].m_sEFIPartitionLetter.c_str() + CString(L":\\");
        CFileDialog fileDlg(TRUE, NULL, NULL,
            OFN_ALLOWMULTISELECT | OFN_EXPLORER,
            L"All Files (*.*)|*.*||", this);
        fileDlg.m_ofn.lpstrTitle = L"Browse the EFI partition";
        fileDlg.m_ofn.lpstrInitialDir = initialDir;
        fileDlg.DoModal();
    }
}

void CEFIMounterDlg::OnClose()
{
    if (!m_bShutdownPerformed) {
        if (m_nCurrentDiskIndex != -1 && !m_disks[m_nCurrentDiskIndex].m_sEFIPartitionLetter.empty()) {
            try {
                SetStatus(L"Unmounting before exit...", true);
                // For simplicity, do this synchronously on close.
                bool updateGuid = m_checkUpdateGuid.GetCheck() == BST_CHECKED;
                Diskpart::RemoveEFIPartitionLetter(m_disks[m_nCurrentDiskIndex], updateGuid);
            }
            catch (...) {}
        }
        m_bShutdownPerformed = true;
        CoUninitialize();
        CDialogEx::OnClose();
    }
}

LRESULT CEFIMounterDlg::OnDiskLoadComplete(WPARAM wParam, LPARAM lParam)
{
    if (wParam) { // Success
        auto* pDisks = (std::vector<CDisk>*)lParam;
        m_disks = *pDisks;
        delete pDisks;

        for (const auto& disk : m_disks) {
            CString item;
            item.Format(L"Disk %d - %s - %s - %llu GB",
                disk.m_nIndex,
                disk.m_sName.c_str(),
                disk.m_sType.c_str(),
                disk.m_ullSizeInGB);
            m_comboDrives.AddString(item);
        }
    }
    else { // Failure
        CString* pErrorMsg = (CString*)lParam;
        MessageBox(*pErrorMsg, L"Error", MB_OK | MB_ICONERROR);
        delete pErrorMsg;
    }

    EnableControls(TRUE);
    if (m_disks.empty()) {
        SetStatus(L"No drives found", false);
        m_btnExplore.EnableWindow(FALSE);
        m_btnMount.EnableWindow(FALSE);
        m_btnUnmount.EnableWindow(FALSE);
        m_checkUpdateGuid.EnableWindow(FALSE);
    }
    else {
        SetStatus(L"Select a drive", false);
        m_comboDrives.SetCurSel(0);
        OnCbnSelchangeDrivesCombo();
    }

    return 0;
}

LRESULT CEFIMounterDlg::OnMountComplete(WPARAM wParam, LPARAM lParam)
{
    EnableControls(TRUE);
    if (wParam) { // Success
        m_nCurrentDiskIndex = static_cast<int>(lParam);
        UpdateUIStateForSelection();

        // Check whether browsing needs to be triggered after mounting
        if (m_bExploreAfterMount) {
            m_bExploreAfterMount = false; // Reset the flag
            CString initialDir = m_disks[m_nCurrentDiskIndex].m_sEFIPartitionLetter.c_str() + CString(L":\\");
            CFileDialog fileDlg(TRUE, NULL, NULL,
                OFN_ALLOWMULTISELECT | OFN_EXPLORER,
                L"All Files (*.*)|*.*||", this);
            fileDlg.m_ofn.lpstrTitle = L"Browse the EFI partition";
            fileDlg.m_ofn.lpstrInitialDir = initialDir;
            fileDlg.DoModal();
        }
    }
    else { // Failure
        m_bExploreAfterMount = false; // Reset the flag too
        CString* pErrorMsg = (CString*)lParam;
        MessageBox(*pErrorMsg, L"Mount Error", MB_OK | MB_ICONERROR);
        delete pErrorMsg;
        if (m_nCurrentDiskIndex != -1) {
            m_disks[m_nCurrentDiskIndex].m_sEFIPartitionLetter.clear();
        }
        UpdateUIStateForSelection();
    }
    return 0;
}

LRESULT CEFIMounterDlg::OnUnmountComplete(WPARAM wParam, LPARAM lParam)
{
    EnableControls(TRUE);
    if (wParam) { // Success
        m_nCurrentDiskIndex = static_cast<int>(lParam);
        UpdateUIStateForSelection();
    }
    else { // Failure
        CString* pErrorMsg = (CString*)lParam;
        MessageBox(*pErrorMsg, L"Unmount Error", MB_OK | MB_ICONERROR);
        delete pErrorMsg;
        UpdateUIStateForSelection(); // Refresh UI, might still show as mounted
    }
    return 0;
}