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