Ventoy2Disk/Ventoy2Disk/Ventoy2Disk.c

   1 /******************************************************************************
   2  * Ventoy2Disk.c
   3  *
   4  * Copyright (c) 2020, 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 "resource.h"
  23 #include "Language.h"
  24 #include "Ventoy2Disk.h"
  25
  26 PHY_DRIVE_INFO *g_PhyDriveList = NULL;
  27 DWORD g_PhyDriveCount = 0;
  28 static int g_FilterRemovable = 0;
  29 int g_FilterUSB = 1;
  30 int g_ForceOperation = 1;
  31 int g_NoNeedInputYes = 0;
  32 int g_WriteImage = 0;
  33
  34 int ParseCmdLineOption(LPSTR lpCmdLine)
  35 {
  36     int i;
  37     char cfgfile[MAX_PATH];
  38
  39     if (lpCmdLine && lpCmdLine[0])
  40     {
  41         Log("CmdLine:<%s>", lpCmdLine);
  42     }
  43
  44     for (i = 0; i < __argc; i++)
  45     {
  46         if (strncmp(__argv[i], "-U", 2) == 0 ||
  47                         strncmp(__argv[i], "-u", 2) == 0)
  48         {
  49             g_FilterUSB = 0;
  50         }
  51         else if (strncmp(__argv[i], "-F", 2) == 0)
  52         {
  53             g_ForceOperation = 1;
  54         }
  55         else if (strncmp(__argv[i], "-Y", 2) == 0 || strncmp(__argv[i], "-y", 2) == 0)
  56         {
  57             g_NoNeedInputYes = 1;
  58         }
  59     }
  60
  61     GetCurrentDirectoryA(sizeof(cfgfile), cfgfile);
  62     strcat_s(cfgfile, sizeof(cfgfile), "\\Ventoy2Disk.ini");
  63
  64     if (0 == GetPrivateProfileIntA("Filter", "USB", 1, cfgfile))
  65     {
  66         g_FilterUSB = 0;
  67     }
  68
  69     if (1 == GetPrivateProfileIntA("Operation", "Force", 0, cfgfile))
  70     {
  71         g_ForceOperation = 1;
  72     }
  73
  74     Log("Control Flag: %d %d %d", g_FilterRemovable, g_FilterUSB, g_ForceOperation);
  75
  76     return 0;
  77 }
  78
  79 static BOOL IsVentoyPhyDrive(int PhyDrive, UINT64 SizeBytes, MBR_HEAD *pMBR, UINT64 *Part2StartSector, UINT64 *GptPart2Attr)
  80 {
  81     int i;
  82     BOOL bRet;
  83     DWORD dwSize;
  84     HANDLE hDrive;
  85     MBR_HEAD MBR;
  86     UINT32 PartStartSector;
  87     UINT32 PartSectorCount;
  88     CHAR PhyDrivePath[128];
  89     CHAR GUIDStr[128];
  90     GUID ZeroGuid = { 0 };
  91         VTOY_GPT_INFO *pGpt = NULL;
  92
  93     safe_sprintf(PhyDrivePath, "\\\\.\\PhysicalDrive%d", PhyDrive);
  94     hDrive = CreateFileA(PhyDrivePath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
  95     Log("Create file Handle:%p %s status:%u", hDrive, PhyDrivePath, LASTERR);
  96
  97     if (hDrive == INVALID_HANDLE_VALUE)
  98     {
  99         return FALSE;
 100     }
 101
 102         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
 103         Log("Read MBR Ret:%u Size:%u code:%u", bRet, dwSize, LASTERR);
 104
 105     if ((!bRet) || (dwSize != sizeof(MBR)))
 106     {
 107                 CHECK_CLOSE_HANDLE(hDrive);
 108         return FALSE;
 109     }
 110
 111     if (MBR.Byte55 != 0x55 || MBR.ByteAA != 0xAA)
 112     {
 113         Log("Byte55 ByteAA not match 0x%x 0x%x", MBR.Byte55, MBR.ByteAA);
 114                 CHECK_CLOSE_HANDLE(hDrive);
 115         return FALSE;
 116     }
 117
 118
 119
 120         if (MBR.PartTbl[0].FsFlag == 0xEE)
 121         {
 122                 pGpt = malloc(sizeof(VTOY_GPT_INFO));
 123                 if (!pGpt)
 124                 {
 125                         CHECK_CLOSE_HANDLE(hDrive);
 126                         return FALSE;
 127                 }
 128
 129                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
 130                 bRet = ReadFile(hDrive, pGpt, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
 131                 CHECK_CLOSE_HANDLE(hDrive);
 132                 if ((!bRet) || (dwSize != sizeof(VTOY_GPT_INFO)))
 133                 {
 134                         Log("Failed to read gpt info %d %u %d", bRet, dwSize, LASTERR);
 135                         return FALSE;
 136                 }
 137
 138                 if (memcmp(pGpt->Head.Signature, "EFI PART", 8))
 139                 {
 140                         Log("Invalid GPT signature");
 141                         return FALSE;
 142                 }
 143
 144         for (i = 0; i < 128; i++)
 145         {
 146             if (memcmp(&(pGpt->PartTbl[i].PartGuid), &ZeroGuid, sizeof(GUID)) == 0)
 147             {
 148                 continue;
 149             }
 150
 151             Log("=========== Disk%d GPT Partition %d ============", PhyDrive, i + 1);
 152
 153             Log("PartTbl.PartType = %s", GUID2String(&pGpt->PartTbl[i].PartType, GUIDStr, sizeof(GUIDStr)));
 154             Log("PartTbl.PartGuid = %s", GUID2String(&pGpt->PartTbl[i].PartGuid, GUIDStr, sizeof(GUIDStr)));
 155             Log("PartTbl.StartLBA = %llu", (ULONGLONG)pGpt->PartTbl[i].StartLBA);
 156             Log("PartTbl.LastLBA = %llu", (ULONGLONG)pGpt->PartTbl[i].LastLBA);
 157             Log("PartTbl.Attribute = 0x%llx", pGpt->PartTbl[i].Attr);
 158             Log("PartTbl.Name = %S", pGpt->PartTbl[i].Name);
 159         }
 160
 161                 if (memcmp(pGpt->PartTbl[1].Name, L"VTOYEFI", 7 * 2))
 162                 {
 163                         if (pGpt->PartTbl[1].Name[0])
 164                         {
 165                                 Log("Invalid ventoy efi part name <%S>", pGpt->PartTbl[1].Name);
 166                         }
 167                         else
 168                         {
 169                                 Log("Invalid ventoy efi part name <null>");
 170                         }
 171
 172                         return FALSE;
 173                 }
 174
 175         if (pGpt->PartTbl[0].StartLBA != 2048)
 176         {
 177             Log("Part1 not match %llu", pGpt->PartTbl[0].StartLBA);
 178             return FALSE;
 179         }
 180
 181         PartSectorCount = VENTOY_EFI_PART_SIZE / 512;
 182
 183         if (pGpt->PartTbl[1].StartLBA != pGpt->PartTbl[0].LastLBA + 1 ||
 184             (UINT32)(pGpt->PartTbl[1].LastLBA + 1 - pGpt->PartTbl[1].StartLBA) != PartSectorCount)
 185         {
 186             Log("Part2 not match [%llu %llu] [%llu %llu]",
 187                 pGpt->PartTbl[0].StartLBA, pGpt->PartTbl[0].LastLBA,
 188                 pGpt->PartTbl[1].StartLBA, pGpt->PartTbl[1].LastLBA);
 189             return FALSE;
 190         }
 191
 192         *GptPart2Attr = pGpt->PartTbl[1].Attr;
 193                 *Part2StartSector = pGpt->PartTbl[1].StartLBA;
 194
 195         memcpy(pMBR, &(pGpt->MBR), sizeof(MBR_HEAD));
 196         }
 197         else
 198         {
 199                 CHECK_CLOSE_HANDLE(hDrive);
 200
 201         for (i = 0; i < 4; i++)
 202         {
 203             Log("=========== Disk%d MBR Partition %d ============", PhyDrive, i + 1);
 204             Log("PartTbl.Active = 0x%x", MBR.PartTbl[i].Active);
 205             Log("PartTbl.FsFlag = 0x%x", MBR.PartTbl[i].FsFlag);
 206             Log("PartTbl.StartSectorId = %u", MBR.PartTbl[i].StartSectorId);
 207             Log("PartTbl.SectorCount = %u", MBR.PartTbl[i].SectorCount);
 208             Log("PartTbl.StartHead = %u", MBR.PartTbl[i].StartHead);
 209             Log("PartTbl.StartSector = %u", MBR.PartTbl[i].StartSector);
 210             Log("PartTbl.StartCylinder = %u", MBR.PartTbl[i].StartCylinder);
 211             Log("PartTbl.EndHead = %u", MBR.PartTbl[i].EndHead);
 212             Log("PartTbl.EndSector = %u", MBR.PartTbl[i].EndSector);
 213             Log("PartTbl.EndCylinder = %u", MBR.PartTbl[i].EndCylinder);
 214         }
 215
 216                 if (MBR.PartTbl[0].StartSectorId != 2048)
 217                 {
 218                         Log("Part1 not match %u", MBR.PartTbl[0].StartSectorId);
 219                         return FALSE;
 220                 }
 221
 222                 PartStartSector = MBR.PartTbl[0].StartSectorId + MBR.PartTbl[0].SectorCount;
 223                 PartSectorCount = VENTOY_EFI_PART_SIZE / 512;
 224
 225                 if (MBR.PartTbl[1].StartSectorId != PartStartSector ||
 226                         MBR.PartTbl[1].SectorCount != PartSectorCount)
 227                 {
 228                         Log("Part2 not match [0x%x 0x%x] [%u %u] [%u %u]",
 229                                 MBR.PartTbl[1].FsFlag, 0xEF,
 230                                 MBR.PartTbl[1].StartSectorId, PartStartSector,
 231                                 MBR.PartTbl[1].SectorCount, PartSectorCount);
 232                         return FALSE;
 233                 }
 234
 235                 if (MBR.PartTbl[0].Active != 0x80 && MBR.PartTbl[1].Active != 0x80)
 236                 {
 237                         Log("Part1 and Part2 are both NOT active 0x%x 0x%x", MBR.PartTbl[0].Active, MBR.PartTbl[1].Active);
 238             if (MBR.PartTbl[2].Active != 0x80 && MBR.PartTbl[3].Active != 0x80)
 239             {
 240                 Log("Part3 and Part4 are both NOT active 0x%x 0x%x", MBR.PartTbl[2].Active, MBR.PartTbl[3].Active);
 241                 //return FALSE;
 242             }
 243                 }
 244
 245                 *Part2StartSector = MBR.PartTbl[1].StartSectorId;
 246
 247         memcpy(pMBR, &MBR, sizeof(MBR_HEAD));
 248         }
 249
 250     Log("PhysicalDrive%d is ventoy disk", PhyDrive);
 251     return TRUE;
 252 }
 253
 254 int GetVolumeClusterSize(char Drive)
 255 {
 256     CHAR Volume[32] = { 0 };
 257     DWORD SectorsPerCluster = 0;
 258     DWORD BytesPerSector = 0;
 259     DWORD NumberOfFreeClusters = 0;
 260     DWORD TotalNumberOfClusters = 0;
 261
 262     sprintf_s(Volume, sizeof(Volume), "%C:\\", Drive);
 263
 264     if (GetDiskFreeSpaceA(Volume, &SectorsPerCluster, &BytesPerSector, &NumberOfFreeClusters, &TotalNumberOfClusters))
 265     {
 266         Log("GetVolumeClusterSize %s SUCCESS %u %u %u", Volume, SectorsPerCluster, BytesPerSector, SectorsPerCluster * BytesPerSector);
 267         return (int)(SectorsPerCluster * BytesPerSector);
 268     }
 269     else
 270     {
 271         Log("GetVolumeClusterSize %s failed err:%u", Volume, LASTERR);
 272     }
 273
 274     return 0;
 275 }
 276
 277 static int GetVentoyFsNameInPhyDrive(PHY_DRIVE_INFO* CurDrive)
 278 {
 279     int i = 0;
 280     UINT64 Offset;
 281
 282     CHAR Volume[128] = { 0 };
 283     CHAR FsName[MAX_PATH] = { 0 };
 284
 285     while (CurDrive->DriveLetters[i])
 286     {
 287         if (GetPhyDriveByLogicalDrive(CurDrive->DriveLetters[i], &Offset) >= 0)
 288         {
 289             if (Offset == SIZE_1MB)
 290             {
 291                 sprintf_s(Volume, sizeof(Volume), "%C:\\", CurDrive->DriveLetters[i]);
 292                 Log("Find the partition 1 logical drive is %s", Volume);
 293                 break;
 294             }
 295         }
 296         i++;
 297     }
 298
 299     sprintf_s(CurDrive->VentoyFsType, sizeof(CurDrive->VentoyFsType), "??");
 300
 301     if (Volume[0])
 302     {
 303         CurDrive->VentoyFsClusterSize = GetVolumeClusterSize(Volume[0]);
 304
 305         if (GetVolumeInformationA(Volume, NULL, 0, NULL, NULL, NULL, FsName, MAX_PATH))
 306         {
 307             if (_stricmp(FsName, "exFAT") == 0)
 308             {
 309                 sprintf_s(CurDrive->VentoyFsType, sizeof(CurDrive->VentoyFsType), "exFAT");
 310             }
 311             else if (_stricmp(FsName, "NTFS") == 0)
 312             {
 313                 sprintf_s(CurDrive->VentoyFsType, sizeof(CurDrive->VentoyFsType), "NTFS");
 314             }
 315             else if (_stricmp(FsName, "FAT") == 0 || _stricmp(FsName, "FAT32") == 0)
 316             {
 317                 sprintf_s(CurDrive->VentoyFsType, sizeof(CurDrive->VentoyFsType), "FAT32");
 318             }
 319             else
 320             {
 321                 sprintf_s(CurDrive->VentoyFsType, sizeof(CurDrive->VentoyFsType), "%s", FsName);
 322             }
 323
 324             Log("GetVentoyFsNameInPhyDrive %d %s <%s> <%s>", CurDrive->PhyDrive, Volume, FsName, CurDrive->VentoyFsType);
 325         }
 326         else
 327         {
 328             Log("GetVolumeInformationA %s failed %u", Volume, LASTERR);
 329         }
 330     }
 331     else
 332     {
 333         Log("GetVentoyFsNameInPhyDrive %s not found", Volume);
 334     }
 335
 336     return 0;
 337 }
 338
 339 static int FilterPhysicalDrive(PHY_DRIVE_INFO *pDriveList, DWORD DriveCount)
 340 {
 341     DWORD i;
 342     DWORD LogDrive;
 343     int Count = 0;
 344     int Letter = 'A';
 345     int Id = 0;
 346     int LetterCount = 0;
 347     UINT64 Part2GPTAttr = 0;
 348         UINT64 Part2StartSector = 0;
 349     PHY_DRIVE_INFO *CurDrive;
 350         MBR_HEAD MBR;
 351     int LogLetter[VENTOY_MAX_PHY_DRIVE];
 352     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 353
 354     for (LogDrive = GetLogicalDrives(); LogDrive > 0; LogDrive >>= 1)
 355     {
 356         if (LogDrive & 0x01)
 357         {
 358             LogLetter[LetterCount] = Letter;
 359                         PhyDriveId[LetterCount] = GetPhyDriveByLogicalDrive(Letter, NULL);
 360
 361             Log("Logical Drive:%C  ===> PhyDrive:%d", LogLetter[LetterCount], PhyDriveId[LetterCount]);
 362             LetterCount++;
 363         }
 364
 365         Letter++;
 366     }
 367
 368     for (i = 0; i < DriveCount; i++)
 369     {
 370         Part2GPTAttr = 0;
 371         CurDrive = pDriveList + i;
 372
 373         CurDrive->Id = -1;
 374         memset(CurDrive->DriveLetters, 0, sizeof(CurDrive->DriveLetters));
 375
 376         if (g_FilterRemovable && (!CurDrive->RemovableMedia))
 377         {
 378             Log("<%s %s> is filtered for not removable.", CurDrive->VendorId, CurDrive->ProductId);
 379             continue;
 380         }
 381
 382         if (g_FilterUSB && CurDrive->BusType != BusTypeUsb)
 383         {
 384             Log("<%s %s> is filtered for not USB type.", CurDrive->VendorId, CurDrive->ProductId);
 385             continue;
 386         }
 387
 388         CurDrive->Id = Id++;
 389
 390         for (Count = 0, Letter = 0; Letter < LetterCount; Letter++)
 391         {
 392             if (PhyDriveId[Letter] == CurDrive->PhyDrive)
 393             {
 394                 if (Count + 1 < sizeof(CurDrive->DriveLetters) / sizeof(CHAR))
 395                 {
 396                     CurDrive->DriveLetters[Count] = LogLetter[Letter];
 397                 }
 398                 Count++;
 399             }
 400         }
 401
 402         if (IsVentoyPhyDrive(CurDrive->PhyDrive, CurDrive->SizeInBytes, &MBR, &Part2StartSector, &Part2GPTAttr))
 403         {
 404             memcpy(&(CurDrive->MBR), &MBR, sizeof(MBR));
 405             CurDrive->PartStyle = (MBR.PartTbl[0].FsFlag == 0xEE) ? 1 : 0;
 406             CurDrive->Part2GPTAttr = Part2GPTAttr;
 407             GetVentoyVerInPhyDrive(CurDrive, Part2StartSector, CurDrive->VentoyVersion, sizeof(CurDrive->VentoyVersion), &(CurDrive->SecureBootSupport));
 408             Log("PhyDrive %d is Ventoy Disk ver:%s SecureBoot:%u", CurDrive->PhyDrive, CurDrive->VentoyVersion, CurDrive->SecureBootSupport);
 409
 410             GetVentoyFsNameInPhyDrive(CurDrive);
 411
 412             if (CurDrive->VentoyVersion[0] == 0)
 413             {
 414                 CurDrive->VentoyVersion[0] = '?';
 415                 Log("Unknown Ventoy Version");
 416             }
 417         }
 418     }
 419
 420     // for safe
 421     for (i = 0; i < DriveCount; i++)
 422     {
 423         CurDrive = pDriveList + i;
 424         if (CurDrive->Id < 0)
 425         {
 426             CurDrive->PhyDrive = 0x00FFFFFF;
 427         }
 428     }
 429
 430     return Id;
 431 }
 432
 433 PHY_DRIVE_INFO * GetPhyDriveInfoById(int Id)
 434 {
 435     DWORD i;
 436     for (i = 0; i < g_PhyDriveCount; i++)
 437     {
 438         if (g_PhyDriveList[i].Id >= 0 && g_PhyDriveList[i].Id == Id)
 439         {
 440             return g_PhyDriveList + i;
 441         }
 442     }
 443
 444     return NULL;
 445 }
 446
 447
 448 PHY_DRIVE_INFO * GetPhyDriveInfoByPhyDrive(int PhyDrive)
 449 {
 450         DWORD i;
 451         for (i = 0; i < g_PhyDriveCount; i++)
 452         {
 453                 if (g_PhyDriveList[i].PhyDrive == PhyDrive)
 454                 {
 455                         return g_PhyDriveList + i;
 456                 }
 457         }
 458
 459         return NULL;
 460 }
 461
 462
 463
 464 int SortPhysicalDrive(PHY_DRIVE_INFO *pDriveList, DWORD DriveCount)
 465 {
 466         DWORD i, j;
 467         BOOL flag;
 468         PHY_DRIVE_INFO TmpDriveInfo;
 469
 470         for (i = 0; i < DriveCount; i++)
 471         {
 472                 for (j = i + 1; j < DriveCount; j++)
 473                 {
 474                         flag = FALSE;
 475
 476                         if (pDriveList[i].BusType == BusTypeUsb && pDriveList[j].BusType == BusTypeUsb)
 477                         {
 478                                 if (pDriveList[i].RemovableMedia == FALSE && pDriveList[j].RemovableMedia == TRUE)
 479                                 {
 480                                         flag = TRUE;
 481                                 }
 482                         }
 483                         else if (pDriveList[j].BusType == BusTypeUsb)
 484                         {
 485                                 flag = TRUE;
 486                         }
 487                         else
 488                         {
 489                                 if (pDriveList[j].PhyDrive < pDriveList[i].PhyDrive)
 490                                 {
 491                                         flag = TRUE;
 492                                 }
 493                         }
 494
 495                         if (flag)
 496                         {
 497                                 memcpy(&TmpDriveInfo, pDriveList + i, sizeof(PHY_DRIVE_INFO));
 498                                 memcpy(pDriveList + i, pDriveList + j, sizeof(PHY_DRIVE_INFO));
 499                                 memcpy(pDriveList + j, &TmpDriveInfo, sizeof(PHY_DRIVE_INFO));
 500                         }
 501                 }
 502         }
 503
 504         return 0;
 505 }
 506
 507 int Ventoy2DiskInit(void)
 508 {
 509     Log("\n===================== Enum All PhyDrives =====================");
 510     g_PhyDriveList = (PHY_DRIVE_INFO *)malloc(sizeof(PHY_DRIVE_INFO)* VENTOY_MAX_PHY_DRIVE);
 511     if (NULL == g_PhyDriveList)
 512     {
 513         Log("Failed to alloc phy drive memory");
 514         return FALSE;
 515     }
 516     memset(g_PhyDriveList, 0, sizeof(PHY_DRIVE_INFO)* VENTOY_MAX_PHY_DRIVE);
 517
 518     GetAllPhysicalDriveInfo(g_PhyDriveList, &g_PhyDriveCount);
 519
 520         SortPhysicalDrive(g_PhyDriveList, g_PhyDriveCount);
 521
 522     FilterPhysicalDrive(g_PhyDriveList, g_PhyDriveCount);
 523
 524     return 0;
 525 }
 526
 527 int Ventoy2DiskDestroy(void)
 528 {
 529     free(g_PhyDriveList);
 530         g_PhyDriveList = NULL;
 531         g_PhyDriveCount = 0;
 532
 533     return 0;
 534 }