Ventoy2Disk/Ventoy2Disk/DiskService_wmsa.c

   1 /******************************************************************************
   2 * DiskService_wsma.c
   3 *
   4 * Copyright (c) 2021, longpanda <admin@ventoy.net>
   5 *
   6 * This program is free software; you can redistribute it and/or
   7 * modify it under the terms of the GNU General Public License as
   8 * published by the Free Software Foundation; either version 3 of the
   9 * License, or (at your option) any later version.
  10 *
  11 * This program is distributed in the hope that it will be useful, but
  12 * WITHOUT ANY WARRANTY; without even the implied warranty of
  13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14 * General Public License for more details.
  15 *
  16 * You should have received a copy of the GNU General Public License
  17 * along with this program; if not, see <http://www.gnu.org/licenses/>.
  18 *
  19 */
  20
  21 #include <Windows.h>
  22 #include <winternl.h>
  23 #include <commctrl.h>
  24 #include <initguid.h>
  25 #include <vds.h>
  26 #include <VersionHelpers.h>
  27 #include "Ventoy2Disk.h"
  28 #include "DiskService.h"
  29
  30 STATIC BOOL IsPowershellExist(void)
  31 {
  32         BOOL ret;
  33
  34         if (!IsWindows8OrGreater())
  35         {
  36                 Log("This is before Windows8 powershell disk not supported.");
  37                 return FALSE;
  38         }
  39
  40         ret = IsFileExist("C:\\Windows\\system32\\WindowsPowerShell\\v1.0\\powershell.exe");
  41         if (!ret)
  42         {
  43                 Log("powershell.exe not exist");
  44         }
  45
  46         return ret;
  47 }
  48
  49 int PSHELL_GetPartitionNumber(int PhyDrive, UINT64 Offset)
  50 {
  51         int partnum = -1;
  52         DWORD i = 0;
  53         DWORD BufLen = 0;
  54         DWORD dwBytes = 0;
  55         BOOL bRet;
  56         HANDLE hDrive;
  57         LONGLONG PartStart;
  58         DRIVE_LAYOUT_INFORMATION_EX *pDriveLayout = NULL;
  59
  60         Log("PSHELL_GetPartitionNumber PhyDrive:%d Offset:%llu", PhyDrive, Offset);
  61
  62         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
  63         if (hDrive == INVALID_HANDLE_VALUE)
  64         {
  65                 return -1;
  66         }
  67
  68         BufLen = (DWORD)(sizeof(PARTITION_INFORMATION_EX)* 256);
  69
  70         pDriveLayout = malloc(BufLen);
  71         if (!pDriveLayout)
  72         {
  73                 goto out;
  74         }
  75         memset(pDriveLayout, 0, BufLen);
  76
  77         bRet = DeviceIoControl(hDrive,
  78                 IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL,
  79                 0,
  80                 pDriveLayout,
  81                 BufLen,
  82                 &dwBytes,
  83                 NULL);
  84         if (!bRet)
  85         {
  86                 Log("Failed to ioctrl get drive layout ex %u", LASTERR);
  87                 goto out;
  88         }
  89
  90         Log("PhyDrive:%d  PartitionStyle=%s  PartitionCount=%u", PhyDrive,
  91                 (pDriveLayout->PartitionStyle == PARTITION_STYLE_MBR) ? "MBR" : "GPT", pDriveLayout->PartitionCount);
  92
  93         for (i = 0; i < pDriveLayout->PartitionCount; i++)
  94         {
  95                 PartStart = pDriveLayout->PartitionEntry[i].StartingOffset.QuadPart;
  96                 if (PartStart == (LONGLONG)Offset)
  97                 {
  98                         Log("[*] [%d] PartitionNumber=%u Offset=%lld Length=%lld ",
  99                                 i,
 100                                 pDriveLayout->PartitionEntry[i].PartitionNumber,
 101                                 pDriveLayout->PartitionEntry[i].StartingOffset.QuadPart,
 102                                 pDriveLayout->PartitionEntry[i].PartitionLength.QuadPart
 103                                 );
 104                         partnum = (int)pDriveLayout->PartitionEntry[i].PartitionNumber;
 105                 }
 106                 else
 107                 {
 108                         Log("[ ] [%d] PartitionNumber=%u Offset=%lld Length=%lld ",
 109                                 i,
 110                                 pDriveLayout->PartitionEntry[i].PartitionNumber,
 111                                 pDriveLayout->PartitionEntry[i].StartingOffset.QuadPart,
 112                                 pDriveLayout->PartitionEntry[i].PartitionLength.QuadPart
 113                                 );
 114                 }
 115         }
 116
 117 out:
 118
 119         CHECK_CLOSE_HANDLE(hDrive);
 120         CHECK_FREE(pDriveLayout);
 121
 122         return partnum;
 123 }
 124
 125
 126 STATIC BOOL PSHELL_CommProc(const char *Cmd)
 127 {
 128         CHAR CmdBuf[4096];
 129         STARTUPINFOA Si;
 130         PROCESS_INFORMATION Pi;
 131
 132         if (!IsPowershellExist())
 133         {
 134                 return FALSE;
 135         }
 136
 137         GetStartupInfoA(&Si);
 138         Si.dwFlags |= STARTF_USESHOWWINDOW;
 139         Si.wShowWindow = SW_HIDE;
 140
 141         sprintf_s(CmdBuf, sizeof(CmdBuf), "C:\\Windows\\system32\\WindowsPowerShell\\v1.0\\powershell.exe -Command \"&{ %s }\"", Cmd);
 142
 143         Log("CreateProcess <%s>", CmdBuf);
 144         CreateProcessA(NULL, CmdBuf, NULL, NULL, FALSE, 0, NULL, NULL, &Si, &Pi);
 145
 146         Log("Wair process ...");
 147         WaitForSingleObject(Pi.hProcess, INFINITE);
 148         Log("Process finished...");
 149
 150         CHECK_CLOSE_HANDLE(Pi.hProcess);
 151         CHECK_CLOSE_HANDLE(Pi.hThread);
 152
 153         return TRUE;
 154 }
 155
 156
 157 BOOL PSHELL_CleanDisk(int DriveIndex)
 158 {
 159         BOOL ret;
 160         CHAR CmdBuf[512];
 161
 162         sprintf_s(CmdBuf, sizeof(CmdBuf), "Clear-Disk -Number %d -RemoveData -RemoveOEM -Confirm:$false", DriveIndex);
 163         ret = PSHELL_CommProc(CmdBuf);
 164         Log("CleanDiskByPowershell<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 165
 166         return ret;
 167 }
 168
 169
 170 BOOL PSHELL_DeleteVtoyEFIPartition(int DriveIndex, UINT64 EfiPartOffset)
 171 {
 172         int Part;
 173         BOOL ret;
 174         CHAR CmdBuf[512];
 175
 176         Part = PSHELL_GetPartitionNumber(DriveIndex, EfiPartOffset);
 177         if (Part < 0)
 178         {
 179                 ret = FALSE;
 180         }
 181         else
 182         {
 183                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Remove-Partition -DiskNumber %d -PartitionNumber %d -Confirm:$false", DriveIndex, Part);
 184                 ret = PSHELL_CommProc(CmdBuf);
 185         }
 186
 187         Log("PSHELL_DeleteVtoyEFIPartition<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 188         return ret;
 189 }
 190
 191
 192 BOOL PSHELL_ChangeVtoyEFI2ESP(int DriveIndex, UINT64 Offset)
 193 {
 194         int Part;
 195         BOOL ret;
 196         CHAR CmdBuf[512];
 197
 198         Part = PSHELL_GetPartitionNumber(DriveIndex, Offset);
 199         if (Part < 0)
 200         {
 201                 ret = FALSE;
 202         }
 203         else
 204         {
 205                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Set-Partition -DiskNumber %d -PartitionNumber %d -gpttype '{C12A7328-F81F-11D2-BA4B-00A0C93EC93B}' -Confirm:$false", DriveIndex, Part);
 206                 ret = PSHELL_CommProc(CmdBuf);
 207         }
 208
 209         Log("PSHELL_ChangeVtoyEFI2ESP<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 210         return ret;
 211 }
 212
 213
 214 BOOL PSHELL_ChangeVtoyEFI2Basic(int DriveIndex, UINT64 Offset)
 215 {
 216         int Part;
 217         BOOL ret;
 218         CHAR CmdBuf[512];
 219
 220         Part = PSHELL_GetPartitionNumber(DriveIndex, Offset);
 221         if (Part < 0)
 222         {
 223                 ret = FALSE;
 224         }
 225         else
 226         {
 227                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Set-Partition -DiskNumber %d -PartitionNumber %d -gpttype '{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}' -Confirm:$false", DriveIndex, Part);
 228                 ret = PSHELL_CommProc(CmdBuf);
 229         }
 230
 231         Log("PSHELL_ChangeVtoyEFI2Basic<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 232         return ret;
 233 }
 234
 235 BOOL PSHELL_ShrinkVolume(int DriveIndex, const char* VolumeGuid, CHAR DriveLetter, UINT64 OldBytes, UINT64 ReduceBytes)
 236 {
 237         int Part;
 238         BOOL ret;
 239         CHAR CmdBuf[512];
 240
 241         (void)VolumeGuid;
 242
 243         Part = PSHELL_GetPartitionNumber(DriveIndex, SIZE_1MB);
 244         if (Part < 0)
 245         {
 246                 ret = FALSE;
 247         }
 248         else
 249         {
 250                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Resize-Partition -DiskNumber %d -PartitionNumber %d -Size %llu -Confirm:$false",
 251                         DriveIndex, Part, OldBytes - ReduceBytes);
 252                 ret = PSHELL_CommProc(CmdBuf);
 253         }
 254
 255         Log("PSHELL_ShrinkVolume<%d> %C: ret:%d (%s)", DriveIndex, DriveLetter, ret, ret ? "SUCCESS" : "FAIL");
 256         return ret;
 257 }