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             else if (_stricmp(opt + 4, "UDF") == 0)
 237             {
 238                 fstype = VTOY_FS_UDF;
 239             }
 240         }
 241     }
 242
 243     if (op < 0 || PhyDrive < 0)
 244     {
 245         Log("[ERROR] Invalid parameters %d %d", op, PhyDrive);
 246         return 1;
 247     }
 248
 249     Log("Ventoy CLI %s PhyDrive:%d %s SecureBoot:%d ReserveSpace:%dMB USBCheck:%u FS:%s NonDest:%d",
 250         op == 0 ? "install" : "update",
 251         PhyDrive, PartStyle ? "GPT" : "MBR",
 252         g_SecureBoot, ReserveMB, USBCheck, GetVentoyFsFmtNameByTypeA(fstype), NonDest
 253         );
 254
 255     if (CLI_GetPhyDriveInfo(PhyDrive, pDrvInfo))
 256     {
 257         Log("[ERROR] Failed to get phydrive%d info", PhyDrive);
 258         return 1;
 259     }
 260
 261     Log("PhyDrive:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 262         pDrvInfo->PhyDrive, GetBusTypeString(pDrvInfo->BusType), pDrvInfo->RemovableMedia,
 263         GetHumanReadableGBSize(pDrvInfo->SizeInBytes), pDrvInfo->SizeInBytes,
 264         pDrvInfo->VendorId, pDrvInfo->ProductId);
 265
 266     if (IsVentoyPhyDrive(PhyDrive, pDrvInfo->SizeInBytes, &MBR, &Part2StartSector, &Part2GPTAttr))
 267     {
 268         memcpy(&(pDrvInfo->MBR), &MBR, sizeof(MBR));
 269         pDrvInfo->PartStyle = (MBR.PartTbl[0].FsFlag == 0xEE) ? 1 : 0;
 270         pDrvInfo->Part2GPTAttr = Part2GPTAttr;
 271         GetVentoyVerInPhyDrive(pDrvInfo, Part2StartSector, pDrvInfo->VentoyVersion, sizeof(pDrvInfo->VentoyVersion), &(pDrvInfo->SecureBootSupport));
 272         Log("PhyDrive %d is Ventoy Disk ver:%s SecureBoot:%u", pDrvInfo->PhyDrive, pDrvInfo->VentoyVersion, pDrvInfo->SecureBootSupport);
 273
 274         GetVentoyFsNameInPhyDrive(pDrvInfo);
 275
 276         if (pDrvInfo->VentoyVersion[0] == 0)
 277         {
 278             pDrvInfo->VentoyVersion[0] = '?';
 279             Log("Unknown Ventoy Version");
 280         }
 281     }
 282
 283     if (op == 0 && NonDest)
 284     {
 285         GetLettersBelongPhyDrive(PhyDrive, pDrvInfo->DriveLetters, sizeof(pDrvInfo->DriveLetters));
 286     }
 287
 288     pCfg->op = op;
 289     pCfg->PartStyle = PartStyle;
 290     pCfg->ReserveMB = ReserveMB;
 291     pCfg->USBCheck = USBCheck;
 292     pCfg->NonDest = NonDest;
 293     pCfg->fstype = fstype;
 294
 295     return 0;
 296 }
 297
 298 static int Ventoy_CLI_NonDestInstall(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG* pCfg)
 299 {
 300     int rc;
 301     int TryId = 1;
 302
 303     Log("Ventoy_CLI_NonDestInstall start ...");
 304
 305     if (pDrvInfo->BytesPerLogicalSector == 4096 && pDrvInfo->BytesPerPhysicalSector == 4096)
 306     {
 307         Log("Ventoy does not support 4k native disk.");
 308         rc = 1;
 309         goto out;
 310     }
 311
 312     if (!PartResizePreCheck(NULL))
 313     {
 314         Log("#### Part Resize PreCheck Failed ####");
 315         rc = 1;
 316         goto out;
 317     }
 318
 319     rc = PartitionResizeForVentoy(pDrvInfo);
 320
 321 out:
 322     Log("Ventoy_CLI_NonDestInstall [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 323
 324     return rc;
 325 }
 326
 327
 328 static int Ventoy_CLI_Install(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG *pCfg)
 329 {
 330     int rc;
 331     int TryId = 1;
 332
 333     Log("Ventoy_CLI_Install start ...");
 334
 335     if (pDrvInfo->BytesPerLogicalSector == 4096 && pDrvInfo->BytesPerPhysicalSector == 4096)
 336     {
 337         Log("Ventoy does not support 4k native disk.");
 338         rc = 1;
 339         goto out;
 340     }
 341
 342     if (pCfg->ReserveMB > 0)
 343     {
 344         CLISetReserveSpace(pCfg->ReserveMB);
 345     }
 346
 347     SetVentoyFsType(pCfg->fstype);
 348
 349     rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 350     if (rc)
 351     {
 352         Log("This time install failed, clean disk by disk, wait 3s and retry...");
 353         DISK_CleanDisk(pDrvInfo->PhyDrive);
 354
 355         Sleep(3000);
 356
 357         Log("Now retry to install...");
 358         rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 359
 360         if (rc)
 361         {
 362             Log("This time install failed, clean disk by diskpart, wait 5s and retry...");
 363             DSPT_CleanDisk(pDrvInfo->PhyDrive);
 364
 365             Sleep(5000);
 366
 367             Log("Now retry to install...");
 368             rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 369         }
 370     }
 371
 372     SetVentoyFsType(VTOY_FS_EXFAT);
 373
 374 out:
 375     Log("Ventoy_CLI_Install [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 376
 377     return rc;
 378 }
 379
 380 static int Ventoy_CLI_Update(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG* pCfg)
 381 {
 382     int rc;
 383     int TryId = 1;
 384
 385     Log("Ventoy_CLI_Update start ...");
 386
 387     rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 388     if (rc)
 389     {
 390         Log("This time update failed, now wait and retry...");
 391         Sleep(4000);
 392
 393         //Try2
 394         Log("Now retry to update...");
 395         rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 396         if (rc)
 397         {
 398             //Try3
 399             Sleep(1000);
 400             Log("Now retry to update...");
 401             rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 402             if (rc)
 403             {
 404                 //Try4 is dangerous ...
 405                 Sleep(3000);
 406                 Log("Now retry to update...");
 407                 rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 408             }
 409         }
 410     }
 411
 412     Log("Ventoy_CLI_Update [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 413
 414     return rc;
 415 }
 416
 417 void CLI_UpdatePercent(int Pos)
 418 {
 419     int Len;
 420     FILE* File = NULL;
 421     CHAR szBuf[128];
 422
 423     Len = (int)sprintf_s(szBuf, sizeof(szBuf), "%d", Pos * 100 / PT_FINISH);
 424     fopen_s(&File, VENTOY_CLI_PERCENT, "w+");
 425     if (File)
 426     {
 427         fwrite(szBuf, 1, Len, File);
 428         fwrite("\n", 1, 1, File);
 429         fclose(File);
 430     }
 431 }
 432
 433 static void CLI_WriteDoneFile(int ret)
 434 {
 435     FILE* File = NULL;
 436
 437     fopen_s(&File, VENTOY_CLI_DONE, "w+");
 438     if (File)
 439     {
 440         if (ret == 0)
 441         {
 442             fwrite("0\n", 1, 2, File);
 443         }
 444         else
 445         {
 446             fwrite("1\n", 1, 2, File);
 447         }
 448         fclose(File);
 449     }
 450 }
 451
 452 PHY_DRIVE_INFO* CLI_PhyDrvInfo(void)
 453 {
 454     return g_CLI_PhyDrvInfo;
 455 }
 456
 457 /*
 458  * Ventoy2Disk.exe VTOYCLI  { /I | /U }  { /Drive:F: | /PhyDrive:1 }  /GPT  /NoSB  /R:4096 /NoUSBCheck
 459  *
 460  */
 461 int VentoyCLIMain(int argc, char** argv)
 462 {
 463     int ret = 1;
 464     PHY_DRIVE_INFO* pDrvInfo = NULL;
 465     CLI_CFG CliCfg;
 466
 467     DeleteFileA(VENTOY_CLI_PERCENT);
 468     DeleteFileA(VENTOY_CLI_DONE);
 469
 470     g_CLI_PhyDrvInfo = pDrvInfo = (PHY_DRIVE_INFO*)malloc(sizeof(PHY_DRIVE_INFO));
 471     if (!pDrvInfo)
 472     {
 473         goto end;
 474     }
 475     memset(pDrvInfo, 0, sizeof(PHY_DRIVE_INFO));
 476
 477     if (CLI_CheckParam(argc, argv, pDrvInfo, &CliCfg))
 478     {
 479         goto end;
 480     }
 481
 482     //Check USB type for install
 483     if (CliCfg.op == 0 && CliCfg.USBCheck)
 484     {
 485         if (pDrvInfo->BusType != BusTypeUsb)
 486         {
 487             Log("[ERROR] PhyDrive %d is NOT USB type",  pDrvInfo->PhyDrive);
 488             goto end;
 489         }
 490     }
 491
 492     if (CliCfg.op == 0)
 493     {
 494         if (CliCfg.NonDest)
 495         {
 496             ret = Ventoy_CLI_NonDestInstall(pDrvInfo, &CliCfg);
 497         }
 498         else
 499         {
 500             AlertSuppressInit();
 501             SetAlertPromptHookEnable(TRUE);
 502             ret = Ventoy_CLI_Install(pDrvInfo, &CliCfg);
 503         }
 504     }
 505     else
 506     {
 507         if (pDrvInfo->VentoyVersion[0] == 0)
 508         {
 509             Log("[ERROR] No Ventoy information detected in PhyDrive %d, so can not do update", pDrvInfo->PhyDrive);
 510             goto end;
 511         }
 512
 513         ret = Ventoy_CLI_Update(pDrvInfo, &CliCfg);
 514     }
 515
 516 end:
 517     CHECK_FREE(pDrvInfo);
 518
 519     CLI_UpdatePercent(PT_FINISH);
 520     CLI_WriteDoneFile(ret);
 521
 522     return ret;
 523 }