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 }CLI_CFG;
  21
  22 BOOL g_CLI_Mode = FALSE;
  23 static int g_CLI_OP;
  24 static int g_CLI_PhyDrive;
  25
  26 static int CLI_GetPhyDriveInfo(int PhyDrive, PHY_DRIVE_INFO* pInfo)
  27 {
  28     BOOL bRet;
  29     DWORD dwBytes;
  30     HANDLE Handle = INVALID_HANDLE_VALUE;
  31     CHAR PhyDrivePath[128];
  32     GET_LENGTH_INFORMATION LengthInfo;
  33     STORAGE_PROPERTY_QUERY Query;
  34     STORAGE_DESCRIPTOR_HEADER DevDescHeader;
  35     STORAGE_DEVICE_DESCRIPTOR* pDevDesc;
  36
  37     safe_sprintf(PhyDrivePath, "\\\\.\\PhysicalDrive%d", PhyDrive);
  38     Handle = CreateFileA(PhyDrivePath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
  39     Log("Create file Handle:%p %s status:%u", Handle, PhyDrivePath, LASTERR);
  40
  41     if (Handle == INVALID_HANDLE_VALUE)
  42     {
  43         return 1;
  44     }
  45
  46     bRet = DeviceIoControl(Handle,
  47         IOCTL_DISK_GET_LENGTH_INFO, NULL,
  48         0,
  49         &LengthInfo,
  50         sizeof(LengthInfo),
  51         &dwBytes,
  52         NULL);
  53     if (!bRet)
  54     {
  55         Log("DeviceIoControl IOCTL_DISK_GET_LENGTH_INFO failed error:%u", LASTERR);
  56         return 1;
  57     }
  58
  59     Log("PHYSICALDRIVE%d size %llu bytes", PhyDrive, (ULONGLONG)LengthInfo.Length.QuadPart);
  60
  61     Query.PropertyId = StorageDeviceProperty;
  62     Query.QueryType = PropertyStandardQuery;
  63
  64     bRet = DeviceIoControl(Handle,
  65         IOCTL_STORAGE_QUERY_PROPERTY,
  66         &Query,
  67         sizeof(Query),
  68         &DevDescHeader,
  69         sizeof(STORAGE_DESCRIPTOR_HEADER),
  70         &dwBytes,
  71         NULL);
  72     if (!bRet)
  73     {
  74         Log("DeviceIoControl1 error:%u dwBytes:%u", LASTERR, dwBytes);
  75         return 1;
  76     }
  77
  78     if (DevDescHeader.Size < sizeof(STORAGE_DEVICE_DESCRIPTOR))
  79     {
  80         Log("Invalid DevDescHeader.Size:%u", DevDescHeader.Size);
  81         return 1;
  82     }
  83
  84     pDevDesc = (STORAGE_DEVICE_DESCRIPTOR*)malloc(DevDescHeader.Size);
  85     if (!pDevDesc)
  86     {
  87         Log("failed to malloc error:%u len:%u", LASTERR, DevDescHeader.Size);
  88         return 1;
  89     }
  90
  91     bRet = DeviceIoControl(Handle,
  92         IOCTL_STORAGE_QUERY_PROPERTY,
  93         &Query,
  94         sizeof(Query),
  95         pDevDesc,
  96         DevDescHeader.Size,
  97         &dwBytes,
  98         NULL);
  99     if (!bRet)
 100     {
 101         Log("DeviceIoControl2 error:%u dwBytes:%u", LASTERR, dwBytes);
 102         free(pDevDesc);
 103         return 1;
 104     }
 105
 106     pInfo->PhyDrive = PhyDrive;
 107     pInfo->SizeInBytes = LengthInfo.Length.QuadPart;
 108     pInfo->DeviceType = pDevDesc->DeviceType;
 109     pInfo->RemovableMedia = pDevDesc->RemovableMedia;
 110     pInfo->BusType = pDevDesc->BusType;
 111
 112     if (pDevDesc->VendorIdOffset)
 113     {
 114         safe_strcpy(pInfo->VendorId, (char*)pDevDesc + pDevDesc->VendorIdOffset);
 115         TrimString(pInfo->VendorId);
 116     }
 117
 118     if (pDevDesc->ProductIdOffset)
 119     {
 120         safe_strcpy(pInfo->ProductId, (char*)pDevDesc + pDevDesc->ProductIdOffset);
 121         TrimString(pInfo->ProductId);
 122     }
 123
 124     if (pDevDesc->ProductRevisionOffset)
 125     {
 126         safe_strcpy(pInfo->ProductRev, (char*)pDevDesc + pDevDesc->ProductRevisionOffset);
 127         TrimString(pInfo->ProductRev);
 128     }
 129
 130     if (pDevDesc->SerialNumberOffset)
 131     {
 132         safe_strcpy(pInfo->SerialNumber, (char*)pDevDesc + pDevDesc->SerialNumberOffset);
 133         TrimString(pInfo->SerialNumber);
 134     }
 135
 136     free(pDevDesc);
 137
 138     CHECK_CLOSE_HANDLE(Handle);
 139
 140     return 0;
 141 }
 142
 143 static int CLI_CheckParam(int argc, char** argv, PHY_DRIVE_INFO* pDrvInfo, CLI_CFG *pCfg)
 144 {
 145     int i;
 146     int op = -1;
 147     char* opt = NULL;
 148     int PhyDrive = -1;
 149     int PartStyle = 0;
 150     int ReserveMB = 0;
 151     BOOL USBCheck = TRUE;
 152     MBR_HEAD MBR;
 153     UINT64 Part2GPTAttr = 0;
 154     UINT64 Part2StartSector = 0;
 155
 156     for (i = 0; i < argc; i++)
 157     {
 158         opt = argv[i];
 159         if (_stricmp(opt, "/I") == 0)
 160         {
 161             op = 0;
 162         }
 163         else if (_stricmp(opt, "/U") == 0)
 164         {
 165             op = 1;
 166         }
 167         else if (_stricmp(opt, "/GPT") == 0)
 168         {
 169             PartStyle = 1;
 170         }
 171         else if (_stricmp(opt, "/NoSB") == 0)
 172         {
 173             g_SecureBoot = FALSE;
 174         }
 175         else if (_stricmp(opt, "/NoUSBCheck") == 0)
 176         {
 177             USBCheck = FALSE;
 178         }
 179         else if (_strnicmp(opt, "/Drive:", 7) == 0)
 180         {
 181             Log("Get PhyDrive by logical drive %C:", opt[7]);
 182             PhyDrive = GetPhyDriveByLogicalDrive(opt[7], NULL);
 183         }
 184         else if (_strnicmp(opt, "/PhyDrive:", 10) == 0)
 185         {
 186             PhyDrive = (int)strtol(opt + 10, NULL, 10);
 187         }
 188         else if (_strnicmp(opt, "/R:", 3) == 0)
 189         {
 190             ReserveMB = (int)strtol(opt + 3, NULL, 10);
 191         }
 192     }
 193
 194     if (op < 0 || PhyDrive < 0)
 195     {
 196         Log("[ERROR] Invalid parameters %d %d", op, PhyDrive);
 197         return 1;
 198     }
 199
 200     Log("Ventoy CLI %s PhyDrive:%d %s SecureBoot:%d ReserveSpace:%dMB USBCheck:%u",
 201         op == 0 ? "install" : "update",
 202         PhyDrive, PartStyle ? "GPT" : "MBR",
 203         g_SecureBoot, ReserveMB, USBCheck
 204         );
 205
 206     if (CLI_GetPhyDriveInfo(PhyDrive, pDrvInfo))
 207     {
 208         Log("[ERROR] Failed to get phydrive%d info", PhyDrive);
 209         return 1;
 210     }
 211
 212     Log("PhyDrive:%d BusType:%-4s Removable:%u Size:%dGB(%llu) Name:%s %s",
 213         pDrvInfo->PhyDrive, GetBusTypeString(pDrvInfo->BusType), pDrvInfo->RemovableMedia,
 214         GetHumanReadableGBSize(pDrvInfo->SizeInBytes), pDrvInfo->SizeInBytes,
 215         pDrvInfo->VendorId, pDrvInfo->ProductId);
 216
 217     if (IsVentoyPhyDrive(PhyDrive, pDrvInfo->SizeInBytes, &MBR, &Part2StartSector, &Part2GPTAttr))
 218     {
 219         memcpy(&(pDrvInfo->MBR), &MBR, sizeof(MBR));
 220         pDrvInfo->PartStyle = (MBR.PartTbl[0].FsFlag == 0xEE) ? 1 : 0;
 221         pDrvInfo->Part2GPTAttr = Part2GPTAttr;
 222         GetVentoyVerInPhyDrive(pDrvInfo, Part2StartSector, pDrvInfo->VentoyVersion, sizeof(pDrvInfo->VentoyVersion), &(pDrvInfo->SecureBootSupport));
 223         Log("PhyDrive %d is Ventoy Disk ver:%s SecureBoot:%u", pDrvInfo->PhyDrive, pDrvInfo->VentoyVersion, pDrvInfo->SecureBootSupport);
 224
 225         GetVentoyFsNameInPhyDrive(pDrvInfo);
 226
 227         if (pDrvInfo->VentoyVersion[0] == 0)
 228         {
 229             pDrvInfo->VentoyVersion[0] = '?';
 230             Log("Unknown Ventoy Version");
 231         }
 232     }
 233
 234     pCfg->op = op;
 235     pCfg->PartStyle = PartStyle;
 236     pCfg->ReserveMB = ReserveMB;
 237     pCfg->USBCheck = USBCheck;
 238
 239     return 0;
 240 }
 241
 242 static int Ventoy_CLI_Install(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG *pCfg)
 243 {
 244     int rc;
 245     int TryId = 1;
 246
 247     Log("Ventoy_CLI_Install start ...");
 248
 249     if (pCfg->ReserveMB > 0)
 250     {
 251         CLISetReserveSpace(pCfg->ReserveMB);
 252     }
 253
 254     rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 255     if (rc)
 256     {
 257         Log("This time install failed, clean disk by disk, wait 3s and retry...");
 258         DISK_CleanDisk(pDrvInfo->PhyDrive);
 259
 260         Sleep(3000);
 261
 262         Log("Now retry to install...");
 263         rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 264
 265         if (rc)
 266         {
 267             Log("This time install failed, clean disk by diskpart, wait 5s and retry...");
 268             DSPT_CleanDisk(pDrvInfo->PhyDrive);
 269
 270             Sleep(5000);
 271
 272             Log("Now retry to install...");
 273             rc = InstallVentoy2PhyDrive(pDrvInfo, pCfg->PartStyle, TryId++);
 274         }
 275     }
 276
 277     Log("Ventoy_CLI_Install [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 278
 279     return rc;
 280 }
 281
 282 static int Ventoy_CLI_Update(PHY_DRIVE_INFO* pDrvInfo, CLI_CFG* pCfg)
 283 {
 284     int rc;
 285     int TryId = 1;
 286
 287     Log("Ventoy_CLI_Update start ...");
 288
 289     rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 290     if (rc)
 291     {
 292         Log("This time update failed, now wait and retry...");
 293         Sleep(4000);
 294
 295         //Try2
 296         Log("Now retry to update...");
 297         rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 298         if (rc)
 299         {
 300             //Try3
 301             Sleep(1000);
 302             Log("Now retry to update...");
 303             rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 304             if (rc)
 305             {
 306                 //Try4 is dangerous ...
 307                 Sleep(3000);
 308                 Log("Now retry to update...");
 309                 rc = UpdateVentoy2PhyDrive(pDrvInfo, TryId++);
 310             }
 311         }
 312     }
 313
 314     Log("Ventoy_CLI_Update [%s]", rc == 0 ? "SUCCESS" : "FAILED");
 315
 316     return rc;
 317 }
 318
 319 void CLI_UpdatePercent(int Pos)
 320 {
 321     int Len;
 322     FILE* File = NULL;
 323     CHAR szBuf[128];
 324
 325     Len = (int)sprintf_s(szBuf, sizeof(szBuf), "%d", Pos * 100 / PT_FINISH);
 326     fopen_s(&File, VENTOY_CLI_PERCENT, "w+");
 327     if (File)
 328     {
 329         fwrite(szBuf, 1, Len, File);
 330         fwrite("\n", 1, 1, File);
 331         fclose(File);
 332     }
 333 }
 334
 335 static void CLI_WriteDoneFile(int ret)
 336 {
 337     FILE* File = NULL;
 338
 339     fopen_s(&File, VENTOY_CLI_DONE, "w+");
 340     if (File)
 341     {
 342         if (ret == 0)
 343         {
 344             fwrite("0\n", 1, 2, File);
 345         }
 346         else
 347         {
 348             fwrite("1\n", 1, 2, File);
 349         }
 350         fclose(File);
 351     }
 352 }
 353
 354 /*
 355  * Ventoy2Disk.exe VTOYCLI  { /I | /U }  { /Drive:F: | /PhyDrive:1 }  /GPT  /NoSB  /R:4096 /NoUSBCheck
 356  *
 357  */
 358 int VentoyCLIMain(int argc, char** argv)
 359 {
 360     int ret = 1;
 361     PHY_DRIVE_INFO* pDrvInfo = NULL;
 362     CLI_CFG CliCfg;
 363
 364     DeleteFileA(VENTOY_CLI_PERCENT);
 365     DeleteFileA(VENTOY_CLI_DONE);
 366
 367     pDrvInfo = (PHY_DRIVE_INFO*)malloc(sizeof(PHY_DRIVE_INFO));
 368     if (!pDrvInfo)
 369     {
 370         goto end;
 371     }
 372     memset(pDrvInfo, 0, sizeof(PHY_DRIVE_INFO));
 373
 374     if (CLI_CheckParam(argc, argv, pDrvInfo, &CliCfg))
 375     {
 376         goto end;
 377     }
 378
 379     //Check USB type for install
 380     if (CliCfg.op == 0 && CliCfg.USBCheck)
 381     {
 382         if (pDrvInfo->BusType != BusTypeUsb)
 383         {
 384             Log("[ERROR] PhyDrive %d is NOT USB type",  pDrvInfo->PhyDrive);
 385             goto end;
 386         }
 387     }
 388
 389     if (CliCfg.op == 0)
 390     {
 391         ret = Ventoy_CLI_Install(pDrvInfo, &CliCfg);
 392     }
 393     else
 394     {
 395         if (pDrvInfo->VentoyVersion[0] == 0)
 396         {
 397             Log("[ERROR] No Ventoy information detected in PhyDrive %d, so can not do update", pDrvInfo->PhyDrive);
 398             goto end;
 399         }
 400
 401         ret = Ventoy_CLI_Update(pDrvInfo, &CliCfg);
 402     }
 403
 404 end:
 405     CHECK_FREE(pDrvInfo);
 406
 407     CLI_UpdatePercent(PT_FINISH);
 408     CLI_WriteDoneFile(ret);
 409
 410     return ret;
 411 }