Ventoy2Disk/Ventoy2Disk/ventoy_cli.c

   1 #include <Windows.h>
   2 #include <Shlobj.h>
   3 #include <tlhelp32.h>
   4 #include <Psapi.h>
   5 #include <commctrl.h>
   6 #include "resource.h"
   7 #include "Language.h"
   8 #include "Ventoy2Disk.h"
   9 #include "DiskService.h"
  10 #include "VentoyJson.h"
  11
  12 extern void CLISetReserveSpace(int MB);
  13
  14 typedef struct CLI_CFG
  15 {
  16     int op;
  17     int PartStyle;
  18     int ReserveMB;
  19     BOOL USBCheck;
  20     BOOL NonDest;
  21     int fstype;
  22 }CLI_CFG;
  23
  24 BOOL g_CLI_Mode = FALSE;
  25 static int g_CLI_OP;
  26 static int g_CLI_PhyDrive;
  27
  28 static PHY_DRIVE_INFO* g_CLI_PhyDrvInfo = NULL;
  29
  30 static int CLI_GetPhyDriveInfo(int PhyDrive, PHY_DRIVE_INFO* pInfo)
  31 {
  32     BOOL bRet;
  33     DWORD dwBytes;
  34     HANDLE Handle = INVALID_HANDLE_VALUE;
  35     CHAR PhyDrivePath[128];
  36     GET_LENGTH_INFORMATION LengthInfo;
  37     STORAGE_PROPERTY_QUERY Query;
  38     STORAGE_DESCRIPTOR_HEADER DevDescHeader;
  39     STORAGE_DEVICE_DESCRIPTOR* pDevDesc;
  40     STORAGE_ACCESS_ALIGNMENT_DESCRIPTOR diskAlignment;
  41
  42     safe_sprintf(PhyDrivePath, "\\\\.\\PhysicalDrive%d", PhyDrive);
  43     Handle = CreateFileA(PhyDrivePath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
  44     Log("Create file Handle:%p %s status:%u", Handle, PhyDrivePath, LASTERR);
  45
  46     if (Handle == INVALID_HANDLE_VALUE)
  47     {
  48         return 1;
  49     }
  50
  51     bRet = DeviceIoControl(Handle,
  52         IOCTL_DISK_GET_LENGTH_INFO, NULL,
  53         0,
  54         &LengthInfo,
  55         sizeof(LengthInfo),
  56         &dwBytes,
  57         NULL);
  58     if (!bRet)
  59     {
  60         Log("DeviceIoControl IOCTL_DISK_GET_LENGTH_INFO failed error:%u", LASTERR);
  61         return 1;
  62     }
  63
  64     Log("PHYSICALDRIVE%d size %llu bytes", PhyDrive, (ULONGLONG)LengthInfo.Length.QuadPart);
  65
  66     Query.PropertyId = StorageDeviceProperty;
  67     Query.QueryType = PropertyStandardQuery;
  68
  69     bRet = DeviceIoControl(Handle,
  70         IOCTL_STORAGE_QUERY_PROPERTY,
  71         &Query,
  72         sizeof(Query),
  73         &DevDescHeader,
  74         sizeof(STORAGE_DESCRIPTOR_HEADER),
  75         &dwBytes,
  76         NULL);
  77     if (!bRet)
  78     {
  79         Log("DeviceIoControl1 error:%u dwBytes:%u", LASTERR, dwBytes);
  80         return 1;
  81     }
  82
  83     if (DevDescHeader.Size < sizeof(STORAGE_DEVICE_DESCRIPTOR))
  84     {
  85         Log("Invalid DevDescHeader.Size:%u", DevDescHeader.Size);
  86         return 1;
  87     }
  88
  89     pDevDesc = (STORAGE_DEVICE_DESCRIPTOR*)malloc(DevDescHeader.Size);
  90     if (!pDevDesc)
  91     {
  92         Log("failed to malloc error:%u len:%u", LASTERR, DevDescHeader.Size);
  93         return 1;
  94     }
  95
  96     bRet = DeviceIoControl(Handle,
  97         IOCTL_STORAGE_QUERY_PROPERTY,
  98         &Query,
  99         sizeof(Query),
 100         pDevDesc,
 101         DevDescHeader.Size,
 102         &dwBytes,
 103         NULL);
 104     if (!bRet)
 105     {
 106         Log("DeviceIoControl2 error:%u dwBytes:%u", LASTERR, dwBytes);
 107         free(pDevDesc);
 108         return 1;
 109     }
 110
 111
 112     memset(&Query, 0, sizeof(STORAGE_PROPERTY_QUERY));
 113     Query.PropertyId = StorageAccessAlignmentProperty;
 114     Query.QueryType = PropertyStandardQuery;
 115     memset(&diskAlignment, 0, sizeof(STORAGE_ACCESS_ALIGNMENT_DESCRIPTOR));
 116
 117     bRet = DeviceIoControl(Handle,
 118         IOCTL_STORAGE_QUERY_PROPERTY,
 119         &Query,
 120         sizeof(STORAGE_PROPERTY_QUERY),
 121         &diskAlignment,
 122         sizeof(STORAGE_ACCESS_ALIGNMENT_DESCRIPTOR),
 123         &dwBytes,
 124         NULL);
 125     if (!bRet)
 126     {
 127         Log("DeviceIoControl3 error:%u dwBytes:%u", LASTERR, dwBytes);
 128     }
 129
 130
 131     pInfo->PhyDrive = PhyDrive;
 132     pInfo->SizeInBytes = LengthInfo.Length.QuadPart;
 133     pInfo->DeviceType = pDevDesc->DeviceType;
 134     pInfo->RemovableMedia = pDevDesc->RemovableMedia;
 135     pInfo->BusType = pDevDesc->BusType;
 136
 137     pInfo->BytesPerLogicalSector = diskAlignment.BytesPerLogicalSector;
 138     pInfo->BytesPerPhysicalSector = diskAlignment.BytesPerPhysicalSector;
 139
 140     if (pDevDesc->VendorIdOffset)
 141     {
 142         safe_strcpy(pInfo->VendorId, (char*)pDevDesc + pDevDesc->VendorIdOffset);
 143         TrimString(pInfo->VendorId);
 144     }
 145
 146     if (pDevDesc->ProductIdOffset)
 147     {
 148         safe_strcpy(pInfo->ProductId, (char*)pDevDesc + pDevDesc->ProductIdOffset);
 149         TrimString(pInfo->ProductId);
 150     }
 151
 152     if (pDevDesc->ProductRevisionOffset)
 153     {
 154         safe_strcpy(pInfo->ProductRev, (char*)pDevDesc + pDevDesc->ProductRevisionOffset);
 155         TrimString(pInfo->ProductRev);
 156     }
 157
 158     if (pDevDesc->SerialNumberOffset)
 159     {
 160         safe_strcpy(pInfo->SerialNumber, (char*)pDevDesc + pDevDesc->SerialNumberOffset);
 161         TrimString(pInfo->SerialNumber);
 162     }
 163
 164     free(pDevDesc);
 165
 166     CHECK_CLOSE_HANDLE(Handle);
 167
 168     return 0;
 169 }
 170
 171 static int CLI_CheckParam(int argc, char** argv, PHY_DRIVE_INFO* pDrvInfo, CLI_CFG *pCfg)
 172 {
 173     int i;
 174     int fstype = VTOY_FS_EXFAT;
 175     int op = -1;
 176     char* opt = NULL;
 177     int PhyDrive = -1;
 178     int PartStyle = 0;
 179     int ReserveMB = 0;
 180     BOOL USBCheck = TRUE;
 181     BOOL NonDest = FALSE;
 182     MBR_HEAD MBR;
 183     UINT64 Part2GPTAttr = 0;
 184     UINT64 Part2StartSector = 0;
 185
 186     for (i = 0; i < argc; i++)
 187     {
 188         opt = argv[i];
 189         if (_stricmp(opt, "/I") == 0)
 190         {
 191             op = 0;
 192         }
 193         else if (_stricmp(opt, "/U") == 0)
 194         {
 195             op = 1;
 196         }
 197         else if (_stricmp(opt, "/GPT") == 0)
 198         {
 199             PartStyle = 1;
 200         }
 201         else if (_stricmp(opt, "/NoSB") == 0)
 202         {
 203             g_SecureBoot = FALSE;
 204         }
 205         else if (_stricmp(opt, "/NoUSBCheck") == 0)
 206         {
 207             USBCheck = FALSE;
 208         }
 209         else if (_stricmp(opt, "/NonDest") == 0)
 210         {
 211             NonDest = TRUE;
 212         }
 213         else if (_strnicmp(opt, "/Drive:", 7) == 0)
 214         {
 215             Log("Get PhyDrive by logical drive %C:", opt[7]);
 216             PhyDrive = GetPhyDriveByLogicalDrive(opt[7], NULL);
 217         }
 218         else if (_strnicmp(opt, "/PhyDrive:", 10) == 0)
 219         {
 220             PhyDrive = (int)strtol(opt + 10, NULL, 10);
 221         }
 222         else if (_strnicmp(opt, "/R:", 3) == 0)
 223         {
 224             ReserveMB = (int)strtol(opt + 3, NULL, 10);
 225         }
 226         else if (_strnicmp(opt, "/FS:", 4) == 0)
 227         {
 228             if (_stricmp(opt + 4, "NTFS") == 0)
 229             {
 230                 fstype = VTOY_FS_NTFS;
 231             }
 232             else if (_stricmp(opt + 4, "FAT32") == 0)
 233             {
 234                 fstype = VTOY_FS_FAT32;
 235             }
 236         }
 237     }
 238
 239     if (op < 0 || PhyDrive < 0)
 240     {
 241         Log("[ERROR] Invalid parameters %d %d", op, PhyDrive);
 242         return 1;
 243     }
 244
 245     Log("Ventoy CLI %s PhyDrive:%d %s SecureBoot:%d ReserveSpace:%dMB USBCheck:%u FS:%s NonDest:%d",
 246         op == 0 ? "install" : "update",
 247         PhyDrive, PartStyle ? "GPT" : "MBR",
 248         g_SecureBoot, ReserveMB, USBCheck, GetVentoyFsFmtNameByTypeA(fstype), NonDest
 249         );
 250
 251     if (CLI_GetPhyDriveInfo(PhyDrive, pDrvInfo))
 252     {
 253         Log("[ERROR] Failed to get phydrive%d info", PhyDrive);
 254         return 1;
 255     }
 256
 257     Log("PhyDrive:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 258         pDrvInfo->PhyDrive, GetBusTypeString(pDrvInfo->BusType), pDrvInfo->RemovableMedia,
 259         GetHumanReadableGBSize(pDrvInfo->SizeInBytes), pDrvInfo->SizeInBytes,
 260         pDrvInfo->VendorId, pDrvInfo->ProductId);
 261
 262     if (IsVentoyPhyDrive(PhyDrive, pDrvInfo->SizeInBytes, &MBR, &Part2StartSector, &Part2GPTAttr))
 263     {
 264         memcpy(&(pDrvInfo->MBR), &MBR, sizeof(MBR));
 265         pDrvInfo->PartStyle = (MBR.PartTbl[0].FsFlag == 0xEE) ? 1 : 0;
 266         pDrvInfo->Part2GPTAttr = Part2GPTAttr;
 267         GetVentoyVerInPhyDrive(pDrvInfo, Part2StartSector, pDrvInfo->VentoyVersion, sizeof(pDrvInfo->VentoyVersion), &(pDrvInfo->SecureBootSupport));
 268         Log("PhyDrive %d is Ventoy Disk ver:%s SecureBoot:%u", pDrvInfo->PhyDrive, pDrvInfo->VentoyVersion, pDrvInfo->SecureBootSupport);
 269
 270         GetVentoyFsNameInPhyDrive(pDrvInfo);
 271
 272         if (pDrvInfo->VentoyVersion[0] == 0)
 273         {
 274             pDrvInfo->VentoyVersion[0] = '?';
 275             Log("Unknown Ventoy Version");
 276         }
 277     }
 278
 279     if (op == 0 && NonDest)
 280     {
 281         GetLettersBelongPhyDrive(PhyDrive, pDrvInfo->DriveLetters, sizeof(pDrvInfo->DriveLetters));
 282     }
 283
 284     pCfg->op = op;
 285     pCfg->PartStyle = PartStyle;
 286     pCfg->ReserveMB = ReserveMB;
 287     pCfg->USBCheck = USBCheck;
 288     pCfg->NonDest = NonDest;
 289     pCfg->fstype = fstype;
 290
 291     return 0;
 292 }
 293
 294 static int Ventoy_CLI_NonDestInstall(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG* pCfg)
 295 {
 296     int rc;
 297     int TryId = 1;
 298
 299     Log("Ventoy_CLI_NonDestInstall start ...");
 300
 301     if (pDrvInfo->BytesPerLogicalSector == 4096 && pDrvInfo->BytesPerPhysicalSector == 4096)
 302     {
 303         Log("Ventoy does not support 4k native disk.");
 304         rc = 1;
 305         goto out;
 306     }
 307
 308     if (!PartResizePreCheck(NULL))
 309     {
 310         Log("#### Part Resize PreCheck Failed ####");
 311         rc = 1;
 312         goto out;
 313     }
 314
 315     rc = PartitionResizeForVentoy(pDrvInfo);
 316
 317 out:
 318     Log("Ventoy_CLI_NonDestInstall [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 319
 320     return rc;
 321 }
 322
 323
 324 static int Ventoy_CLI_Install(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG *pCfg)
 325 {
 326     int rc;
 327     int TryId = 1;
 328
 329     Log("Ventoy_CLI_Install start ...");
 330
 331     if (pDrvInfo->BytesPerLogicalSector == 4096 && pDrvInfo->BytesPerPhysicalSector == 4096)
 332     {
 333         Log("Ventoy does not support 4k native disk.");
 334         rc = 1;
 335         goto out;
 336     }
 337
 338     if (pCfg->ReserveMB > 0)
 339     {
 340         CLISetReserveSpace(pCfg->ReserveMB);
 341     }
 342
 343     SetVentoyFsType(pCfg->fstype);
 344
 345     rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 346     if (rc)
 347     {
 348         Log("This time install failed, clean disk by disk, wait 3s and retry...");
 349         DISK_CleanDisk(pDrvInfo->PhyDrive);
 350
 351         Sleep(3000);
 352
 353         Log("Now retry to install...");
 354         rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 355
 356         if (rc)
 357         {
 358             Log("This time install failed, clean disk by diskpart, wait 5s and retry...");
 359             DSPT_CleanDisk(pDrvInfo->PhyDrive);
 360
 361             Sleep(5000);
 362
 363             Log("Now retry to install...");
 364             rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 365         }
 366     }
 367
 368     SetVentoyFsType(VTOY_FS_EXFAT);
 369
 370 out:
 371     Log("Ventoy_CLI_Install [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 372
 373     return rc;
 374 }
 375
 376 static int Ventoy_CLI_Update(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG* pCfg)
 377 {
 378     int rc;
 379     int TryId = 1;
 380
 381     Log("Ventoy_CLI_Update start ...");
 382
 383     rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 384     if (rc)
 385     {
 386         Log("This time update failed, now wait and retry...");
 387         Sleep(4000);
 388
 389         //Try2
 390         Log("Now retry to update...");
 391         rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 392         if (rc)
 393         {
 394             //Try3
 395             Sleep(1000);
 396             Log("Now retry to update...");
 397             rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 398             if (rc)
 399             {
 400                 //Try4 is dangerous ...
 401                 Sleep(3000);
 402                 Log("Now retry to update...");
 403                 rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 404             }
 405         }
 406     }
 407
 408     Log("Ventoy_CLI_Update [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 409
 410     return rc;
 411 }
 412
 413 void CLI_UpdatePercent(int Pos)
 414 {
 415     int Len;
 416     FILE* File = NULL;
 417     CHAR szBuf[128];
 418
 419     Len = (int)sprintf_s(szBuf, sizeof(szBuf), "%d", Pos * 100 / PT_FINISH);
 420     fopen_s(&File, VENTOY_CLI_PERCENT, "w+");
 421     if (File)
 422     {
 423         fwrite(szBuf, 1, Len, File);
 424         fwrite("\n", 1, 1, File);
 425         fclose(File);
 426     }
 427 }
 428
 429 static void CLI_WriteDoneFile(int ret)
 430 {
 431     FILE* File = NULL;
 432
 433     fopen_s(&File, VENTOY_CLI_DONE, "w+");
 434     if (File)
 435     {
 436         if (ret == 0)
 437         {
 438             fwrite("0\n", 1, 2, File);
 439         }
 440         else
 441         {
 442             fwrite("1\n", 1, 2, File);
 443         }
 444         fclose(File);
 445     }
 446 }
 447
 448 PHY_DRIVE_INFO* CLI_PhyDrvInfo(void)
 449 {
 450     return g_CLI_PhyDrvInfo;
 451 }
 452
 453 /*
 454  * Ventoy2Disk.exe VTOYCLI  { /I | /U }  { /Drive:F: | /PhyDrive:1 }  /GPT  /NoSB  /R:4096 /NoUSBCheck
 455  *
 456  */
 457 int VentoyCLIMain(int argc, char** argv)
 458 {
 459     int ret = 1;
 460     PHY_DRIVE_INFO* pDrvInfo = NULL;
 461     CLI_CFG CliCfg;
 462
 463     DeleteFileA(VENTOY_CLI_PERCENT);
 464     DeleteFileA(VENTOY_CLI_DONE);
 465
 466     g_CLI_PhyDrvInfo = pDrvInfo = (PHY_DRIVE_INFO*)malloc(sizeof(PHY_DRIVE_INFO));
 467     if (!pDrvInfo)
 468     {
 469         goto end;
 470     }
 471     memset(pDrvInfo, 0, sizeof(PHY_DRIVE_INFO));
 472
 473     if (CLI_CheckParam(argc, argv, pDrvInfo, &CliCfg))
 474     {
 475         goto end;
 476     }
 477
 478     //Check USB type for install
 479     if (CliCfg.op == 0 && CliCfg.USBCheck)
 480     {
 481         if (pDrvInfo->BusType != BusTypeUsb)
 482         {
 483             Log("[ERROR] PhyDrive %d is NOT USB type",  pDrvInfo->PhyDrive);
 484             goto end;
 485         }
 486     }
 487
 488     if (CliCfg.op == 0)
 489     {
 490         if (CliCfg.NonDest)
 491         {
 492             ret = Ventoy_CLI_NonDestInstall(pDrvInfo, &CliCfg);
 493         }
 494         else
 495         {
 496             AlertSuppressInit();
 497             SetAlertPromptHookEnable(TRUE);
 498             ret = Ventoy_CLI_Install(pDrvInfo, &CliCfg);
 499         }
 500     }
 501     else
 502     {
 503         if (pDrvInfo->VentoyVersion[0] == 0)
 504         {
 505             Log("[ERROR] No Ventoy information detected in PhyDrive %d, so can not do update", pDrvInfo->PhyDrive);
 506             goto end;
 507         }
 508
 509         ret = Ventoy_CLI_Update(pDrvInfo, &CliCfg);
 510     }
 511
 512 end:
 513     CHECK_FREE(pDrvInfo);
 514
 515     CLI_UpdatePercent(PT_FINISH);
 516     CLI_WriteDoneFile(ret);
 517
 518     return ret;
 519 }