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         ret = IsFileExist("%s\\system32\\WindowsPowerShell\\v1.0\\powershell.exe", DISK_GetWindowsDir());
  35         if (!ret)
  36         {
  37                 Log("powershell.exe not exist");
  38         }
  39
  40         return ret;
  41 }
  42
  43 int PSHELL_GetPartitionNumber(int PhyDrive, UINT64 Offset)
  44 {
  45         int partnum = -1;
  46         DWORD i = 0;
  47         DWORD BufLen = 0;
  48         DWORD dwBytes = 0;
  49         BOOL bRet;
  50         HANDLE hDrive;
  51         LONGLONG PartStart;
  52         DRIVE_LAYOUT_INFORMATION_EX *pDriveLayout = NULL;
  53
  54         Log("PSHELL_GetPartitionNumber PhyDrive:%d Offset:%llu", PhyDrive, Offset);
  55
  56         hDrive = GetPhysicalHandle(PhyDrive, FALSE, FALSE, FALSE);
  57         if (hDrive == INVALID_HANDLE_VALUE)
  58         {
  59                 return -1;
  60         }
  61
  62         BufLen = (DWORD)(sizeof(PARTITION_INFORMATION_EX)* 256);
  63
  64         pDriveLayout = malloc(BufLen);
  65         if (!pDriveLayout)
  66         {
  67                 goto out;
  68         }
  69         memset(pDriveLayout, 0, BufLen);
  70
  71         bRet = DeviceIoControl(hDrive,
  72                 IOCTL_DISK_GET_DRIVE_LAYOUT_EX, NULL,
  73                 0,
  74                 pDriveLayout,
  75                 BufLen,
  76                 &dwBytes,
  77                 NULL);
  78         if (!bRet)
  79         {
  80                 Log("Failed to ioctrl get drive layout ex %u", LASTERR);
  81                 goto out;
  82         }
  83
  84         Log("PhyDrive:%d  PartitionStyle=%s  PartitionCount=%u", PhyDrive,
  85                 (pDriveLayout->PartitionStyle == PARTITION_STYLE_MBR) ? "MBR" : "GPT", pDriveLayout->PartitionCount);
  86
  87         for (i = 0; i < pDriveLayout->PartitionCount; i++)
  88         {
  89                 PartStart = pDriveLayout->PartitionEntry[i].StartingOffset.QuadPart;
  90                 if (PartStart == (LONGLONG)Offset)
  91                 {
  92                         Log("[*] [%d] PartitionNumber=%u Offset=%lld Length=%lld ",
  93                                 i,
  94                                 pDriveLayout->PartitionEntry[i].PartitionNumber,
  95                                 pDriveLayout->PartitionEntry[i].StartingOffset.QuadPart,
  96                                 pDriveLayout->PartitionEntry[i].PartitionLength.QuadPart
  97                                 );
  98                         partnum = (int)pDriveLayout->PartitionEntry[i].PartitionNumber;
  99                 }
 100                 else
 101                 {
 102                         Log("[ ] [%d] PartitionNumber=%u Offset=%lld Length=%lld ",
 103                                 i,
 104                                 pDriveLayout->PartitionEntry[i].PartitionNumber,
 105                                 pDriveLayout->PartitionEntry[i].StartingOffset.QuadPart,
 106                                 pDriveLayout->PartitionEntry[i].PartitionLength.QuadPart
 107                                 );
 108                 }
 109         }
 110
 111 out:
 112
 113         CHECK_CLOSE_HANDLE(hDrive);
 114         CHECK_FREE(pDriveLayout);
 115
 116         return partnum;
 117 }
 118
 119
 120 STATIC BOOL PSHELL_CommProc(const char *Cmd)
 121 {
 122         CHAR CmdBuf[4096];
 123         STARTUPINFOA Si;
 124         PROCESS_INFORMATION Pi;
 125
 126         if (!IsPowershellExist())
 127         {
 128                 return FALSE;
 129         }
 130
 131         GetStartupInfoA(&Si);
 132         Si.dwFlags |= STARTF_USESHOWWINDOW;
 133         Si.wShowWindow = SW_HIDE;
 134
 135         sprintf_s(CmdBuf, sizeof(CmdBuf), "C:\\Windows\\system32\\WindowsPowerShell\\v1.0\\powershell.exe -Command \"&{ %s }\"", Cmd);
 136
 137         Log("CreateProcess <%s>", CmdBuf);
 138         CreateProcessA(NULL, CmdBuf, NULL, NULL, FALSE, 0, NULL, NULL, &Si, &Pi);
 139
 140         Log("Wair process ...");
 141         WaitForSingleObject(Pi.hProcess, INFINITE);
 142         Log("Process finished...");
 143
 144         CHECK_CLOSE_HANDLE(Pi.hProcess);
 145         CHECK_CLOSE_HANDLE(Pi.hThread);
 146
 147         return TRUE;
 148 }
 149
 150
 151 BOOL PSHELL_CleanDisk(int DriveIndex)
 152 {
 153         BOOL ret;
 154         CHAR CmdBuf[512];
 155
 156         sprintf_s(CmdBuf, sizeof(CmdBuf), "Clear-Disk -Number %d -RemoveData -RemoveOEM -Confirm:$false", DriveIndex);
 157         ret = PSHELL_CommProc(CmdBuf);
 158         Log("CleanDiskByPowershell<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 159
 160         return ret;
 161 }
 162
 163
 164 BOOL PSHELL_DeleteVtoyEFIPartition(int DriveIndex, UINT64 EfiPartOffset)
 165 {
 166         int Part;
 167         BOOL ret;
 168         CHAR CmdBuf[512];
 169
 170         Part = PSHELL_GetPartitionNumber(DriveIndex, EfiPartOffset);
 171         if (Part < 0)
 172         {
 173                 ret = FALSE;
 174         }
 175         else
 176         {
 177                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Remove-Partition -DiskNumber %d -PartitionNumber %d -Confirm:$false", DriveIndex, Part);
 178                 ret = PSHELL_CommProc(CmdBuf);
 179         }
 180
 181         Log("PSHELL_DeleteVtoyEFIPartition<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 182         return ret;
 183 }
 184
 185
 186 BOOL PSHELL_ChangeVtoyEFI2ESP(int DriveIndex, UINT64 Offset)
 187 {
 188         int Part;
 189         BOOL ret;
 190         CHAR CmdBuf[512];
 191
 192         Part = PSHELL_GetPartitionNumber(DriveIndex, Offset);
 193         if (Part < 0)
 194         {
 195                 ret = FALSE;
 196         }
 197         else
 198         {
 199                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Set-Partition -DiskNumber %d -PartitionNumber %d -gpttype '{C12A7328-F81F-11D2-BA4B-00A0C93EC93B}' -Confirm:$false", DriveIndex, Part);
 200                 ret = PSHELL_CommProc(CmdBuf);
 201         }
 202
 203         Log("PSHELL_ChangeVtoyEFI2ESP<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 204         return ret;
 205 }
 206
 207
 208 BOOL PSHELL_ChangeVtoyEFI2Basic(int DriveIndex, UINT64 Offset)
 209 {
 210         int Part;
 211         BOOL ret;
 212         CHAR CmdBuf[512];
 213
 214         Part = PSHELL_GetPartitionNumber(DriveIndex, Offset);
 215         if (Part < 0)
 216         {
 217                 ret = FALSE;
 218         }
 219         else
 220         {
 221                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Set-Partition -DiskNumber %d -PartitionNumber %d -gpttype '{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}' -Confirm:$false", DriveIndex, Part);
 222                 ret = PSHELL_CommProc(CmdBuf);
 223         }
 224
 225         Log("PSHELL_ChangeVtoyEFI2Basic<%d> ret:%d (%s)", DriveIndex, ret, ret ? "SUCCESS" : "FAIL");
 226         return ret;
 227 }
 228
 229 BOOL PSHELL_ShrinkVolume(int DriveIndex, const char* VolumeGuid, CHAR DriveLetter, UINT64 OldBytes, UINT64 ReduceBytes)
 230 {
 231         int Part;
 232         BOOL ret;
 233         CHAR CmdBuf[512];
 234
 235         (void)VolumeGuid;
 236
 237         Part = PSHELL_GetPartitionNumber(DriveIndex, SIZE_1MB);
 238         if (Part < 0)
 239         {
 240                 ret = FALSE;
 241         }
 242         else
 243         {
 244                 sprintf_s(CmdBuf, sizeof(CmdBuf), "Resize-Partition -DiskNumber %d -PartitionNumber %d -Size %llu -Confirm:$false",
 245                         DriveIndex, Part, OldBytes - ReduceBytes);
 246                 ret = PSHELL_CommProc(CmdBuf);
 247         }
 248
 249         Log("PSHELL_ShrinkVolume<%d> %C: ret:%d (%s)", DriveIndex, DriveLetter, ret, ret ? "SUCCESS" : "FAIL");
 250         return ret;
 251 }
 252
 253 BOOL PSHELL_FormatVolume(char DriveLetter, int fs, DWORD ClusterSize)
 254 {
 255         BOOL ret;
 256         const char* fsname = NULL;
 257         CHAR CmdBuf[512];
 258         CHAR FsName[128];
 259
 260         fsname = GetVentoyFsFmtNameByTypeA(fs);
 261
 262         if (ClusterSize > 0)
 263         {
 264                 sprintf_s(CmdBuf, sizeof(CmdBuf),
 265                         "format-volume -DriveLetter %C -FileSystem %s -AllocationUnitSize %u -Force -NewFileSystemLabel Ventoy",
 266                         DriveLetter, fsname, ClusterSize);
 267         }
 268         else
 269         {
 270                 sprintf_s(CmdBuf, sizeof(CmdBuf),
 271                         "format-volume -DriveLetter %C -FileSystem %s -Force -NewFileSystemLabel Ventoy",
 272                         DriveLetter, fsname);
 273         }
 274
 275         ret = PSHELL_CommProc(CmdBuf);
 276         Log("PSHELL_FormatVolume %C: ret:%d (%s)", DriveLetter, ret, ret ? "SUCCESS" : "FAIL");
 277         if (!ret)
 278         {
 279                 return FALSE;
 280         }
 281
 282
 283         sprintf_s(CmdBuf, sizeof(CmdBuf), "%C:\\", DriveLetter);
 284         GetVolumeInformationA(CmdBuf, NULL, 0, NULL, NULL, NULL, FsName, sizeof(FsName));
 285         VentoyStringToUpper(FsName);
 286
 287         Log("New fs name after run PSHELL:<%s>", FsName);
 288
 289         if (strcmp(FsName, fsname) == 0)
 290         {
 291                 Log("PSHELL_FormatVolume <%C:> SUCCESS", DriveLetter);
 292                 return TRUE;
 293         }
 294         else
 295         {
 296                 Log("PSHELL_FormatVolume <%C:> FAILED", DriveLetter);
 297                 return FALSE;
 298         }
 299 }