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 static int g_FilterUSB = 1;
  30 int g_ForceOperation = 1;
  31
  32 int ParseCmdLineOption(LPSTR lpCmdLine)
  33 {
  34     int i;
  35     char cfgfile[MAX_PATH];
  36
  37     if (lpCmdLine && lpCmdLine[0])
  38     {
  39         Log("CmdLine:<%s>", lpCmdLine);
  40     }
  41
  42     for (i = 0; i < __argc; i++)
  43     {
  44         if (strncmp(__argv[i], "-U", 2) == 0 ||
  45                         strncmp(__argv[i], "-u", 2) == 0)
  46         {
  47             g_FilterUSB = 0;
  48         }
  49         else if (strncmp(__argv[i], "-F", 2) == 0)
  50         {
  51             g_ForceOperation = 1;
  52         }
  53     }
  54
  55     GetCurrentDirectoryA(sizeof(cfgfile), cfgfile);
  56     strcat_s(cfgfile, sizeof(cfgfile), "\\Ventoy2Disk.ini");
  57
  58     if (0 == GetPrivateProfileIntA("Filter", "USB", 1, cfgfile))
  59     {
  60         g_FilterUSB = 0;
  61     }
  62
  63     if (1 == GetPrivateProfileIntA("Operation", "Force", 0, cfgfile))
  64     {
  65         g_ForceOperation = 1;
  66     }
  67
  68     Log("Control Flag: %d %d %d", g_FilterRemovable, g_FilterUSB, g_ForceOperation);
  69
  70     return 0;
  71 }
  72
  73 static BOOL IsVentoyPhyDrive(int PhyDrive, UINT64 SizeBytes)
  74 {
  75     int i;
  76     BOOL bRet;
  77     DWORD dwSize;
  78     HANDLE hDrive;
  79     MBR_HEAD MBR;
  80     UINT32 PartStartSector;
  81     UINT32 PartSectorCount;
  82     CHAR PhyDrivePath[128];
  83
  84     safe_sprintf(PhyDrivePath, "\\\\.\\PhysicalDrive%d", PhyDrive);
  85     hDrive = CreateFileA(PhyDrivePath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
  86     Log("Create file Handle:%p %s status:%u", hDrive, PhyDrivePath, LASTERR);
  87
  88     if (hDrive == INVALID_HANDLE_VALUE)
  89     {
  90         return FALSE;
  91     }
  92
  93     bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
  94     CHECK_CLOSE_HANDLE(hDrive);
  95
  96     Log("Read MBR Ret:%u Size:%u code:%u", bRet, dwSize, LASTERR);
  97
  98     if ((!bRet) || (dwSize != sizeof(MBR)))
  99     {
 100         return FALSE;
 101     }
 102
 103     if (MBR.Byte55 != 0x55 || MBR.ByteAA != 0xAA)
 104     {
 105         Log("Byte55 ByteAA not match 0x%x 0x%x", MBR.Byte55, MBR.ByteAA);
 106         return FALSE;
 107     }
 108
 109     for (i = 0; i < 4; i++)
 110     {
 111         Log("=========== Partition Table %d ============", i + 1);
 112         Log("PartTbl.Active = 0x%x", MBR.PartTbl[i].Active);
 113         Log("PartTbl.FsFlag = 0x%x", MBR.PartTbl[i].FsFlag);
 114         Log("PartTbl.StartSectorId = %u", MBR.PartTbl[i].StartSectorId);
 115         Log("PartTbl.SectorCount = %u", MBR.PartTbl[i].SectorCount);
 116         Log("PartTbl.StartHead = %u", MBR.PartTbl[i].StartHead);
 117         Log("PartTbl.StartSector = %u", MBR.PartTbl[i].StartSector);
 118         Log("PartTbl.StartCylinder = %u", MBR.PartTbl[i].StartCylinder);
 119         Log("PartTbl.EndHead = %u", MBR.PartTbl[i].EndHead);
 120         Log("PartTbl.EndSector = %u", MBR.PartTbl[i].EndSector);
 121         Log("PartTbl.EndCylinder = %u", MBR.PartTbl[i].EndCylinder);
 122     }
 123
 124     if (MBR.PartTbl[2].SectorCount > 0 || MBR.PartTbl[3].SectorCount > 0)
 125     {
 126         Log("Part3 Part4 are not empty");
 127         return FALSE;
 128     }
 129
 130     PartStartSector = 2048;
 131     PartSectorCount = (UINT32)((SizeBytes - VENTOY_EFI_PART_SIZE - SIZE_1MB) / 512);
 132
 133     if (MBR.PartTbl[0].StartSectorId != PartStartSector ||
 134         MBR.PartTbl[0].SectorCount != PartSectorCount)
 135     {
 136         Log("Part1 not match %u %u", PartStartSector, PartSectorCount);
 137         return FALSE;
 138     }
 139
 140     PartStartSector = (UINT32)((SizeBytes - VENTOY_EFI_PART_SIZE) / 512);
 141     PartSectorCount = VENTOY_EFI_PART_SIZE / 512;
 142
 143     if (MBR.PartTbl[1].FsFlag != 0xEF ||
 144         MBR.PartTbl[1].StartSectorId != PartStartSector ||
 145         MBR.PartTbl[1].SectorCount != PartSectorCount)
 146     {
 147         Log("Part2 not match %u %u", PartStartSector, PartSectorCount);
 148         return FALSE;
 149     }
 150
 151     if (MBR.PartTbl[0].Active != 0x80 && MBR.PartTbl[1].Active != 0x80)
 152     {
 153         Log("Part1 and Part2 are both NOT active 0x%x 0x%x", MBR.PartTbl[0].Active, MBR.PartTbl[1].Active);
 154         return FALSE;
 155     }
 156
 157     Log("PhysicalDrive%d is ventoy disk", PhyDrive);
 158     return TRUE;
 159 }
 160
 161
 162 static int FilterPhysicalDrive(PHY_DRIVE_INFO *pDriveList, DWORD DriveCount)
 163 {
 164     DWORD i;
 165     DWORD LogDrive;
 166     int Letter = 'A';
 167     int Id = 0;
 168     int LetterCount = 0;
 169     PHY_DRIVE_INFO *CurDrive;
 170     int LogLetter[VENTOY_MAX_PHY_DRIVE];
 171     int PhyDriveId[VENTOY_MAX_PHY_DRIVE];
 172
 173     for (LogDrive = GetLogicalDrives(); LogDrive > 0; LogDrive >>= 1)
 174     {
 175         if (LogDrive & 0x01)
 176         {
 177             LogLetter[LetterCount] = Letter;
 178             PhyDriveId[LetterCount] = GetPhyDriveByLogicalDrive(Letter);
 179
 180             Log("Logical Drive:%C  ===> PhyDrive:%d", LogLetter[LetterCount], PhyDriveId[LetterCount]);
 181             LetterCount++;
 182         }
 183
 184         Letter++;
 185     }
 186
 187     for (i = 0; i < DriveCount; i++)
 188     {
 189         CurDrive = pDriveList + i;
 190
 191         CurDrive->Id = -1;
 192         CurDrive->FirstDriveLetter = -1;
 193
 194         // Too big for MBR
 195         if (CurDrive->SizeInBytes > 2199023255552ULL)
 196         {
 197             Log("<%s %s> is filtered for too big for MBR.", CurDrive->VendorId, CurDrive->ProductId);
 198             continue;
 199         }
 200
 201         if (g_FilterRemovable && (!CurDrive->RemovableMedia))
 202         {
 203             Log("<%s %s> is filtered for not removable.", CurDrive->VendorId, CurDrive->ProductId);
 204             continue;
 205         }
 206
 207         if (g_FilterUSB && CurDrive->BusType != BusTypeUsb)
 208         {
 209             Log("<%s %s> is filtered for not USB type.", CurDrive->VendorId, CurDrive->ProductId);
 210             continue;
 211         }
 212
 213         CurDrive->Id = Id++;
 214
 215         for (Letter = 0; Letter < LetterCount; Letter++)
 216         {
 217             if (PhyDriveId[Letter] == CurDrive->PhyDrive)
 218             {
 219                 CurDrive->FirstDriveLetter = LogLetter[Letter];
 220                 break;
 221             }
 222         }
 223
 224         if (IsVentoyPhyDrive(CurDrive->PhyDrive, CurDrive->SizeInBytes))
 225         {
 226             GetVentoyVerInPhyDrive(CurDrive, CurDrive->VentoyVersion, sizeof(CurDrive->VentoyVersion));
 227         }
 228     }
 229
 230     // for safe
 231     for (i = 0; i < DriveCount; i++)
 232     {
 233         CurDrive = pDriveList + i;
 234         if (CurDrive->Id < 0)
 235         {
 236             CurDrive->PhyDrive = 0x00FFFFFF;
 237         }
 238     }
 239
 240     return Id;
 241 }
 242
 243 PHY_DRIVE_INFO * GetPhyDriveInfoById(int Id)
 244 {
 245     DWORD i;
 246     for (i = 0; i < g_PhyDriveCount; i++)
 247     {
 248         if (g_PhyDriveList[i].Id >= 0 && g_PhyDriveList[i].Id == Id)
 249         {
 250             return g_PhyDriveList + i;
 251         }
 252     }
 253
 254     return NULL;
 255 }
 256
 257 int SortPhysicalDrive(PHY_DRIVE_INFO *pDriveList, DWORD DriveCount)
 258 {
 259         DWORD i, j;
 260         PHY_DRIVE_INFO TmpDriveInfo;
 261
 262         for (i = 0; i < DriveCount; i++)
 263         {
 264                 for (j = i + 1; j < DriveCount; j++)
 265                 {
 266                         if (pDriveList[i].BusType == BusTypeUsb && pDriveList[j].BusType == BusTypeUsb)
 267                         {
 268                                 if (pDriveList[i].RemovableMedia == FALSE && pDriveList[j].RemovableMedia == TRUE)
 269                                 {
 270                                         memcpy(&TmpDriveInfo, pDriveList + i, sizeof(PHY_DRIVE_INFO));
 271                                         memcpy(pDriveList + i, pDriveList + j, sizeof(PHY_DRIVE_INFO));
 272                                         memcpy(pDriveList + j, &TmpDriveInfo, sizeof(PHY_DRIVE_INFO));
 273                                 }
 274                         }
 275                 }
 276         }
 277
 278         return 0;
 279 }
 280
 281 int Ventoy2DiskInit(void)
 282 {
 283     g_PhyDriveList = (PHY_DRIVE_INFO *)malloc(sizeof(PHY_DRIVE_INFO)* VENTOY_MAX_PHY_DRIVE);
 284     if (NULL == g_PhyDriveList)
 285     {
 286         Log("Failed to alloc phy drive memory");
 287         return FALSE;
 288     }
 289     memset(g_PhyDriveList, 0, sizeof(PHY_DRIVE_INFO)* VENTOY_MAX_PHY_DRIVE);
 290
 291     GetAllPhysicalDriveInfo(g_PhyDriveList, &g_PhyDriveCount);
 292
 293         SortPhysicalDrive(g_PhyDriveList, g_PhyDriveCount);
 294
 295     FilterPhysicalDrive(g_PhyDriveList, g_PhyDriveCount);
 296
 297     return 0;
 298 }
 299
 300 int Ventoy2DiskDestroy(void)
 301 {
 302     free(g_PhyDriveList);
 303     return 0;
 304 }