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