EFIMounter/Diskpart.cpp at main · nuloperrito/EFIMounter · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
#include "pch.h"
#include "Diskpart.h"
#include <vector>
#include <fstream>
#include <stdexcept>
#include <windows.h>

void Diskpart::Run(const std::vector<std::wstring>& commands)
{
    WCHAR tempPath[MAX_PATH];
    WCHAR scriptPath[MAX_PATH];

    if (GetTempPathW(MAX_PATH, tempPath) == 0) {
        throw std::runtime_error("Could not get temporary path.");
    }
    if (GetTempFileNameW(tempPath, L"DP", 0, scriptPath) == 0) {
        throw std::runtime_error("Could not create temporary script file.");
    }

    std::wofstream scriptFile(scriptPath);
    if (!scriptFile) {
        throw std::runtime_error("Could not open temporary script file for writing.");
    }
    for (const auto& cmd : commands) {
        scriptFile << cmd << std::endl;
    }
    scriptFile.close();

    WCHAR diskpartPath[MAX_PATH];
    GetSystemDirectoryW(diskpartPath, MAX_PATH);
    wcscat_s(diskpartPath, MAX_PATH, L"\\diskpart.exe");

    std::wstring params = L"/s \"" + std::wstring(scriptPath) + L"\"";

    SHELLEXECUTEINFOW sei = { sizeof(sei) };
    sei.fMask = SEE_MASK_NOCLOSEPROCESS | SEE_MASK_FLAG_NO_UI;
    sei.lpVerb = L"runas"; // Request elevation
    sei.lpFile = diskpartPath;
    sei.lpParameters = params.c_str();
    sei.nShow = SW_HIDE;

    if (!ShellExecuteExW(&sei) || sei.hProcess == NULL) {
        DeleteFileW(scriptPath);
        throw std::runtime_error("Failed to execute diskpart.exe. Administrator privileges may be required.");
    }

    WaitForSingleObject(sei.hProcess, INFINITE);

    DWORD exitCode = 0;
    GetExitCodeProcess(sei.hProcess, &exitCode);
    CloseHandle(sei.hProcess);
    DeleteFileW(scriptPath);

    if (exitCode != 0) {
        throw std::runtime_error("Diskpart process returned an error.");
    }
}


void Diskpart::AssignLetterToEFIPartition(const CDisk& disk, bool setGuidForGpt)
{
    if (disk.m_nEfiPartitionIndex <= 0) {
        throw std::runtime_error("Invalid EFI partition index for the selected disk.");
    }
    std::vector<std::wstring> commands;
    commands.push_back(L"select disk " + std::to_wstring(disk.m_nIndex));
    commands.push_back(L"select partition " + std::to_wstring(disk.m_nEfiPartitionIndex));
    if (setGuidForGpt) {
        commands.push_back(L"set id=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7 override");
    }
    commands.push_back(L"assign letter=" + disk.m_sEFIPartitionLetter);
    commands.push_back(L"exit");
    Run(commands);
}

void Diskpart::RemoveEFIPartitionLetter(const CDisk& disk, bool setGuidForGpt)
{
    if (disk.m_nEfiPartitionIndex <= 0) {
        throw std::runtime_error("Invalid EFI partition index for the selected disk.");
    }
    std::vector<std::wstring> commands;
    commands.push_back(L"select disk " + std::to_wstring(disk.m_nIndex));
    commands.push_back(L"select partition " + std::to_wstring(disk.m_nEfiPartitionIndex));
    if (setGuidForGpt) {
        commands.push_back(L"set id=C12A7328-F81F-11D2-BA4B-00A0C93EC93B override");
    }
    commands.push_back(L"remove letter=" + disk.m_sEFIPartitionLetter);
    commands.push_back(L"exit");
    Run(commands);
}