vtoyjump/vtoyjump/vtoyjump.c

   1 /******************************************************************************
   2 * vtoyjump.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 <stdio.h>
  22 #include <stdlib.h>
  23 #include <string.h>
  24 #include <Windows.h>
  25 #include <virtdisk.h>
  26 #include <winioctl.h>
  27 #include <VersionHelpers.h>
  28 #include "vtoyjump.h"
  29 #include "fat_filelib.h"
  30
  31 static ventoy_os_param g_os_param;
  32 static ventoy_windows_data g_windows_data;
  33 static UINT8 g_os_param_reserved[32];
  34 static BOOL g_64bit_system = FALSE;
  35 static ventoy_guid g_ventoy_guid = VENTOY_GUID;
  36
  37 void Log(const char *Fmt, ...)
  38 {
  39         va_list Arg;
  40         int Len = 0;
  41         FILE *File = NULL;
  42         SYSTEMTIME Sys;
  43         char szBuf[1024];
  44
  45         GetLocalTime(&Sys);
  46         Len += sprintf_s(szBuf, sizeof(szBuf),
  47                 "[%4d/%02d/%02d %02d:%02d:%02d.%03d] ",
  48                 Sys.wYear, Sys.wMonth, Sys.wDay,
  49                 Sys.wHour, Sys.wMinute, Sys.wSecond,
  50                 Sys.wMilliseconds);
  51
  52         va_start(Arg, Fmt);
  53         Len += vsnprintf_s(szBuf + Len, sizeof(szBuf)-Len, sizeof(szBuf)-Len, Fmt, Arg);
  54         va_end(Arg);
  55
  56         fopen_s(&File, "ventoy.log", "a+");
  57         if (File)
  58         {
  59                 fwrite(szBuf, 1, Len, File);
  60                 fwrite("\n", 1, 1, File);
  61                 fclose(File);
  62         }
  63 }
  64
  65
  66 static int LoadNtDriver(const char *DrvBinPath)
  67 {
  68         int i;
  69         int rc = 0;
  70         BOOL Ret;
  71         DWORD Status;
  72         SC_HANDLE hServiceMgr;
  73         SC_HANDLE hService;
  74         char name[256] = { 0 };
  75
  76         for (i = (int)strlen(DrvBinPath) - 1; i >= 0; i--)
  77         {
  78                 if (DrvBinPath[i] == '\\' || DrvBinPath[i] == '/')
  79                 {
  80                         sprintf_s(name, sizeof(name), "%s", DrvBinPath + i + 1);
  81                         break;
  82                 }
  83         }
  84
  85         Log("Load NT driver: %s %s", DrvBinPath, name);
  86
  87         hServiceMgr = OpenSCManagerA(NULL, NULL, SC_MANAGER_ALL_ACCESS);
  88         if (hServiceMgr == NULL)
  89         {
  90                 Log("OpenSCManager failed Error:%u", GetLastError());
  91                 return 1;
  92         }
  93
  94         Log("OpenSCManager OK");
  95
  96         hService = CreateServiceA(hServiceMgr,
  97                 name,
  98                 name,
  99                 SERVICE_ALL_ACCESS,
 100                 SERVICE_KERNEL_DRIVER,
 101                 SERVICE_DEMAND_START,
 102                 SERVICE_ERROR_NORMAL,
 103                 DrvBinPath,
 104                 NULL, NULL, NULL, NULL, NULL);
 105         if (hService == NULL)
 106         {
 107                 Status = GetLastError();
 108                 if (Status != ERROR_IO_PENDING && Status != ERROR_SERVICE_EXISTS)
 109                 {
 110                         Log("CreateService failed v %u", Status);
 111                         CloseServiceHandle(hServiceMgr);
 112                         return 1;
 113                 }
 114
 115                 hService = OpenServiceA(hServiceMgr, name, SERVICE_ALL_ACCESS);
 116                 if (hService == NULL)
 117                 {
 118                         Log("OpenService failed %u", Status);
 119                         CloseServiceHandle(hServiceMgr);
 120                         return 1;
 121                 }
 122         }
 123
 124         Log("CreateService imdisk OK");
 125
 126         Ret = StartServiceA(hService, 0, NULL);
 127         if (Ret)
 128         {
 129                 Log("StartService OK");
 130         }
 131         else
 132         {
 133                 Status = GetLastError();
 134                 if (Status == ERROR_SERVICE_ALREADY_RUNNING)
 135                 {
 136                         rc = 0;
 137                 }
 138                 else
 139                 {
 140                         Log("StartService error  %u", Status);
 141                         rc = 1;
 142                 }
 143         }
 144
 145         CloseServiceHandle(hService);
 146         CloseServiceHandle(hServiceMgr);
 147
 148         Log("Load NT driver %s", rc ? "failed" : "success");
 149
 150         return rc;
 151 }
 152
 153 static int ReadWholeFile2Buf(const char *Fullpath, void **Data, DWORD *Size)
 154 {
 155         int rc = 1;
 156         DWORD FileSize;
 157         DWORD dwSize;
 158         HANDLE Handle;
 159         BYTE *Buffer = NULL;
 160
 161         Log("ReadWholeFile2Buf <%s>", Fullpath);
 162
 163         Handle = CreateFileA(Fullpath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 164         if (Handle == INVALID_HANDLE_VALUE)
 165         {
 166                 Log("Could not open the file<%s>, error:%u", Fullpath, GetLastError());
 167                 goto End;
 168         }
 169
 170         FileSize = SetFilePointer(Handle, 0, NULL, FILE_END);
 171
 172         Buffer = malloc(FileSize);
 173         if (!Buffer)
 174         {
 175                 Log("Failed to alloc memory size:%u", FileSize);
 176                 goto End;
 177         }
 178
 179         SetFilePointer(Handle, 0, NULL, FILE_BEGIN);
 180         if (!ReadFile(Handle, Buffer, FileSize, &dwSize, NULL))
 181         {
 182                 Log("ReadFile failed, dwSize:%u  error:%u", dwSize, GetLastError());
 183                 goto End;
 184         }
 185
 186         *Data = Buffer;
 187         *Size = FileSize;
 188
 189         Log("Success read file size:%u", FileSize);
 190
 191         rc = 0;
 192
 193 End:
 194         SAFE_CLOSE_HANDLE(Handle);
 195
 196         return rc;
 197 }
 198
 199 static BOOL CheckPeHead(BYTE *Head)
 200 {
 201         UINT32 PeOffset;
 202
 203         if (Head[0] != 'M' || Head[1] != 'Z')
 204         {
 205                 return FALSE;
 206         }
 207
 208         PeOffset = *(UINT32 *)(Head + 60);
 209         if (*(UINT32 *)(Head + PeOffset) != 0x00004550)
 210         {
 211                 return FALSE;
 212         }
 213
 214         return TRUE;
 215 }
 216
 217 static BOOL IsPe64(BYTE *buffer)
 218 {
 219         DWORD pe_off;
 220
 221         if (!CheckPeHead(buffer))
 222         {
 223                 return FALSE;
 224         }
 225
 226         pe_off = *(UINT32 *)(buffer + 60);
 227         if (*(UINT16 *)(buffer + pe_off + 24) == 0x020b)
 228         {
 229                 return TRUE;
 230         }
 231
 232         return FALSE;
 233 }
 234
 235
 236 static BOOL CheckOsParam(ventoy_os_param *param)
 237 {
 238         UINT32 i;
 239         BYTE Sum = 0;
 240
 241         if (memcmp(&param->guid, &g_ventoy_guid, sizeof(ventoy_guid)))
 242         {
 243                 return FALSE;
 244         }
 245
 246         for (i = 0; i < sizeof(ventoy_os_param); i++)
 247         {
 248                 Sum += *((BYTE *)param + i);
 249         }
 250
 251         if (Sum)
 252         {
 253                 return FALSE;
 254         }
 255
 256         if (param->vtoy_img_location_addr % 4096)
 257         {
 258                 return FALSE;
 259         }
 260
 261         return TRUE;
 262 }
 263
 264 static int SaveBuffer2File(const char *Fullpath, void *Buffer, DWORD Length)
 265 {
 266         int rc = 1;
 267         DWORD dwSize;
 268         HANDLE Handle;
 269
 270         Log("SaveBuffer2File <%s> len:%u", Fullpath, Length);
 271
 272         Handle = CreateFileA(Fullpath, GENERIC_READ | GENERIC_WRITE,
 273                 FILE_SHARE_READ | FILE_SHARE_WRITE, 0, CREATE_NEW, 0, 0);
 274         if (Handle == INVALID_HANDLE_VALUE)
 275         {
 276                 Log("Could not create new file, error:%u", GetLastError());
 277                 goto End;
 278         }
 279
 280         WriteFile(Handle, Buffer, Length, &dwSize, NULL);
 281
 282         rc = 0;
 283
 284 End:
 285         SAFE_CLOSE_HANDLE(Handle);
 286
 287         return rc;
 288 }
 289
 290 static int IsUTF8Encode(const char *src)
 291 {
 292     int i;
 293     const UCHAR *Byte = (const UCHAR *)src;
 294
 295     for (i = 0; i < MAX_PATH && Byte[i]; i++)
 296     {
 297         if (Byte[i] > 127)
 298         {
 299             return 1;
 300         }
 301     }
 302
 303     return 0;
 304 }
 305
 306 static int Utf8ToUtf16(const char* src, WCHAR * dst)
 307 {
 308     int size = MultiByteToWideChar(CP_UTF8, 0, src, -1, dst, 0);
 309     return MultiByteToWideChar(CP_UTF8, 0, src, -1, dst, size + 1);
 310 }
 311
 312 static BOOL IsDirExist(const char *Fmt, ...)
 313 {
 314     va_list Arg;
 315     DWORD Attr;
 316     int UTF8 = 0;
 317     CHAR FilePathA[MAX_PATH];
 318     WCHAR FilePathW[MAX_PATH];
 319
 320     va_start(Arg, Fmt);
 321     vsnprintf_s(FilePathA, sizeof(FilePathA), sizeof(FilePathA), Fmt, Arg);
 322     va_end(Arg);
 323
 324     UTF8 = IsUTF8Encode(FilePathA);
 325
 326     if (UTF8)
 327     {
 328         Utf8ToUtf16(FilePathA, FilePathW);
 329         Attr = GetFileAttributesW(FilePathW);
 330     }
 331     else
 332     {
 333         Attr = GetFileAttributesA(FilePathA);
 334     }
 335
 336     if (Attr != INVALID_FILE_ATTRIBUTES && (Attr & FILE_ATTRIBUTE_DIRECTORY))
 337     {
 338         return TRUE;
 339     }
 340
 341     return FALSE;
 342 }
 343
 344 static BOOL IsFileExist(const char *Fmt, ...)
 345 {
 346         va_list Arg;
 347         HANDLE hFile;
 348         DWORD Attr;
 349     BOOL bRet = FALSE;
 350     int UTF8 = 0;
 351         CHAR FilePathA[MAX_PATH];
 352         WCHAR FilePathW[MAX_PATH];
 353
 354         va_start(Arg, Fmt);
 355         vsnprintf_s(FilePathA, sizeof(FilePathA), sizeof(FilePathA), Fmt, Arg);
 356         va_end(Arg);
 357
 358     UTF8 = IsUTF8Encode(FilePathA);
 359
 360     if (UTF8)
 361     {
 362         Utf8ToUtf16(FilePathA, FilePathW);
 363         hFile = CreateFileW(FilePathW, FILE_READ_EA, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
 364     }
 365     else
 366     {
 367         hFile = CreateFileA(FilePathA, FILE_READ_EA, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
 368     }
 369         if (INVALID_HANDLE_VALUE == hFile)
 370         {
 371         goto out;
 372         }
 373
 374         CloseHandle(hFile);
 375
 376     if (UTF8)
 377     {
 378         Attr = GetFileAttributesW(FilePathW);
 379     }
 380     else
 381     {
 382         Attr = GetFileAttributesA(FilePathA);
 383     }
 384
 385     if (Attr & FILE_ATTRIBUTE_DIRECTORY)
 386     {
 387         goto out;
 388     }
 389
 390     bRet = TRUE;
 391
 392 out:
 393     Log("File <%s> %s", FilePathA, (bRet ? "exist" : "NOT exist"));
 394     return bRet;
 395 }
 396
 397 static int GetPhyDiskUUID(const char LogicalDrive, UINT8 *UUID, DISK_EXTENT *DiskExtent)
 398 {
 399         BOOL Ret;
 400         DWORD dwSize;
 401         HANDLE Handle;
 402         VOLUME_DISK_EXTENTS DiskExtents;
 403         CHAR PhyPath[128];
 404         UINT8 SectorBuf[512];
 405
 406         Log("GetPhyDiskUUID %C", LogicalDrive);
 407
 408         sprintf_s(PhyPath, sizeof(PhyPath), "\\\\.\\%C:", LogicalDrive);
 409         Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 410         if (Handle == INVALID_HANDLE_VALUE)
 411         {
 412                 Log("Could not open the disk<%s>, error:%u", PhyPath, GetLastError());
 413                 return 1;
 414         }
 415
 416         Ret = DeviceIoControl(Handle,
 417                 IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 418                 NULL,
 419                 0,
 420                 &DiskExtents,
 421                 (DWORD)(sizeof(DiskExtents)),
 422                 (LPDWORD)&dwSize,
 423                 NULL);
 424         if (!Ret || DiskExtents.NumberOfDiskExtents == 0)
 425         {
 426                 Log("DeviceIoControl IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS failed, error:%u", GetLastError());
 427                 CloseHandle(Handle);
 428                 return 1;
 429         }
 430         CloseHandle(Handle);
 431
 432         memcpy(DiskExtent, DiskExtents.Extents, sizeof(DiskExtent));
 433         Log("%C: is in PhysicalDrive%d ", LogicalDrive, DiskExtents.Extents[0].DiskNumber);
 434
 435         sprintf_s(PhyPath, sizeof(PhyPath), "\\\\.\\PhysicalDrive%d", DiskExtents.Extents[0].DiskNumber);
 436         Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 437         if (Handle == INVALID_HANDLE_VALUE)
 438         {
 439                 Log("Could not open the disk<%s>, error:%u", PhyPath, GetLastError());
 440                 return 1;
 441         }
 442
 443         if (!ReadFile(Handle, SectorBuf, sizeof(SectorBuf), &dwSize, NULL))
 444         {
 445                 Log("ReadFile failed, dwSize:%u  error:%u", dwSize, GetLastError());
 446                 CloseHandle(Handle);
 447                 return 1;
 448         }
 449
 450         memcpy(UUID, SectorBuf + 0x180, 16);
 451         CloseHandle(Handle);
 452         return 0;
 453 }
 454
 455 static int VentoyMountAnywhere(HANDLE Handle)
 456 {
 457     DWORD Status;
 458     ATTACH_VIRTUAL_DISK_PARAMETERS AttachParameters;
 459
 460     Log("VentoyMountAnywhere");
 461
 462     memset(&AttachParameters, 0, sizeof(AttachParameters));
 463     AttachParameters.Version = ATTACH_VIRTUAL_DISK_VERSION_1;
 464
 465     Status = AttachVirtualDisk(Handle, NULL, ATTACH_VIRTUAL_DISK_FLAG_READ_ONLY | ATTACH_VIRTUAL_DISK_FLAG_PERMANENT_LIFETIME, 0, &AttachParameters, NULL);
 466     if (Status != ERROR_SUCCESS)
 467     {
 468         Log("Failed to attach virtual disk ErrorCode:%u", Status);
 469         return 1;
 470     }
 471
 472     return 0;
 473 }
 474
 475 int VentoyMountY(HANDLE Handle)
 476 {
 477     int  i;
 478     BOOL  bRet = FALSE;
 479     DWORD Status;
 480     DWORD physicalDriveNameSize;
 481     CHAR *Pos = NULL;
 482     WCHAR physicalDriveName[MAX_PATH];
 483     CHAR physicalDriveNameA[MAX_PATH];
 484     CHAR cdromDriveName[MAX_PATH];
 485     ATTACH_VIRTUAL_DISK_PARAMETERS AttachParameters;
 486
 487     Log("VentoyMountY");
 488
 489     memset(&AttachParameters, 0, sizeof(AttachParameters));
 490     AttachParameters.Version = ATTACH_VIRTUAL_DISK_VERSION_1;
 491
 492     Status = AttachVirtualDisk(Handle, NULL, ATTACH_VIRTUAL_DISK_FLAG_READ_ONLY | ATTACH_VIRTUAL_DISK_FLAG_NO_DRIVE_LETTER | ATTACH_VIRTUAL_DISK_FLAG_PERMANENT_LIFETIME, 0, &AttachParameters, NULL);
 493     if (Status != ERROR_SUCCESS)
 494     {
 495         Log("Failed to attach virtual disk ErrorCode:%u", Status);
 496         return 1;
 497     }
 498
 499     memset(physicalDriveName, 0, sizeof(physicalDriveName));
 500     memset(physicalDriveNameA, 0, sizeof(physicalDriveNameA));
 501
 502     physicalDriveNameSize = MAX_PATH;
 503     Status = GetVirtualDiskPhysicalPath(Handle, &physicalDriveNameSize, physicalDriveName);
 504     if (Status != ERROR_SUCCESS)
 505     {
 506         Log("Failed GetVirtualDiskPhysicalPath ErrorCode:%u", Status);
 507         return 1;
 508     }
 509
 510     for (i = 0; physicalDriveName[i]; i++)
 511     {
 512         physicalDriveNameA[i] = toupper((CHAR)(physicalDriveName[i]));
 513     }
 514
 515     Log("physicalDriveNameA=<%s>", physicalDriveNameA);
 516
 517     Pos = strstr(physicalDriveNameA, "CDROM");
 518     if (!Pos)
 519     {
 520         Log("Not cdrom phy drive");
 521         return 1;
 522     }
 523
 524     sprintf_s(cdromDriveName, sizeof(cdromDriveName), "\\Device\\%s", Pos);
 525     Log("cdromDriveName=<%s>", cdromDriveName);
 526
 527     for (i = 0; i < 3 && (bRet == FALSE); i++)
 528     {
 529         Sleep(1000);
 530         bRet = DefineDosDeviceA(DDD_RAW_TARGET_PATH, "Y:", cdromDriveName);
 531         Log("DefineDosDeviceA %s", bRet ? "success" : "failed");
 532     }
 533
 534     return bRet ? 0 : 1;
 535 }
 536
 537 static BOOL VentoyNeedMountY(const char *IsoPath)
 538 {
 539     /* TBD */
 540     return FALSE;
 541 }
 542
 543 static int VentoyAttachVirtualDisk(HANDLE Handle, const char *IsoPath)
 544 {
 545     int DriveYFree;
 546     DWORD Drives;
 547
 548     Drives = GetLogicalDrives();
 549     if ((1 << 24) & Drives)
 550     {
 551         Log("Y: is occupied");
 552         DriveYFree = 0;
 553     }
 554     else
 555     {
 556         Log("Y: is free now");
 557         DriveYFree = 1;
 558     }
 559
 560     if (DriveYFree && VentoyNeedMountY(IsoPath))
 561     {
 562         return VentoyMountY(Handle);
 563     }
 564     else
 565     {
 566         return VentoyMountAnywhere(Handle);
 567     }
 568 }
 569
 570 int VentoyMountISOByAPI(const char *IsoPath)
 571 {
 572         HANDLE Handle;
 573         DWORD Status;
 574         WCHAR wFilePath[512] = { 0 };
 575         VIRTUAL_STORAGE_TYPE StorageType;
 576         OPEN_VIRTUAL_DISK_PARAMETERS OpenParameters;
 577
 578         Log("VentoyMountISOByAPI <%s>", IsoPath);
 579
 580     if (IsUTF8Encode(IsoPath))
 581     {
 582         MultiByteToWideChar(CP_UTF8, 0, IsoPath, (int)strlen(IsoPath), wFilePath, (int)(sizeof(wFilePath) / sizeof(WCHAR)));
 583     }
 584     else
 585     {
 586         MultiByteToWideChar(CP_ACP, 0, IsoPath, (int)strlen(IsoPath), wFilePath, (int)(sizeof(wFilePath) / sizeof(WCHAR)));
 587     }
 588
 589         memset(&StorageType, 0, sizeof(StorageType));
 590         memset(&OpenParameters, 0, sizeof(OpenParameters));
 591
 592         OpenParameters.Version = OPEN_VIRTUAL_DISK_VERSION_1;
 593
 594         Status = OpenVirtualDisk(&StorageType, wFilePath, VIRTUAL_DISK_ACCESS_READ, 0, &OpenParameters, &Handle);
 595         if (Status != ERROR_SUCCESS)
 596         {
 597                 if (ERROR_VIRTDISK_PROVIDER_NOT_FOUND == Status)
 598                 {
 599                         Log("VirtualDisk for ISO file is not supported in current system");
 600                 }
 601                 else
 602                 {
 603                         Log("Failed to open virtual disk ErrorCode:%u", Status);
 604                 }
 605                 return 1;
 606         }
 607
 608         Log("OpenVirtualDisk success");
 609
 610     Status = VentoyAttachVirtualDisk(Handle, IsoPath);
 611         if (Status != ERROR_SUCCESS)
 612         {
 613                 Log("Failed to attach virtual disk ErrorCode:%u", Status);
 614                 CloseHandle(Handle);
 615                 return 1;
 616         }
 617
 618     Log("VentoyAttachVirtualDisk success");
 619
 620         CloseHandle(Handle);
 621         return 0;
 622 }
 623
 624
 625 static HANDLE g_FatPhyDrive;
 626 static UINT64 g_Part2StartSec;
 627
 628 static int CopyFileFromFatDisk(const CHAR* SrcFile, const CHAR *DstFile)
 629 {
 630         int rc = 1;
 631         int size = 0;
 632         char *buf = NULL;
 633         void *flfile = NULL;
 634
 635         Log("CopyFileFromFatDisk (%s)==>(%s)", SrcFile, DstFile);
 636
 637         flfile = fl_fopen(SrcFile, "rb");
 638         if (flfile)
 639         {
 640                 fl_fseek(flfile, 0, SEEK_END);
 641                 size = (int)fl_ftell(flfile);
 642                 fl_fseek(flfile, 0, SEEK_SET);
 643
 644                 buf = (char *)malloc(size);
 645                 if (buf)
 646                 {
 647                         fl_fread(buf, 1, size, flfile);
 648
 649                         rc = 0;
 650                         SaveBuffer2File(DstFile, buf, size);
 651                         free(buf);
 652                 }
 653
 654                 fl_fclose(flfile);
 655         }
 656
 657         return rc;
 658 }
 659
 660 static int VentoyFatDiskRead(uint32 Sector, uint8 *Buffer, uint32 SectorCount)
 661 {
 662         DWORD dwSize;
 663         BOOL bRet;
 664         DWORD ReadSize;
 665         LARGE_INTEGER liCurrentPosition;
 666
 667         liCurrentPosition.QuadPart = Sector + g_Part2StartSec;
 668         liCurrentPosition.QuadPart *= 512;
 669         SetFilePointerEx(g_FatPhyDrive, liCurrentPosition, &liCurrentPosition, FILE_BEGIN);
 670
 671         ReadSize = (DWORD)(SectorCount * 512);
 672
 673         bRet = ReadFile(g_FatPhyDrive, Buffer, ReadSize, &dwSize, NULL);
 674         if (bRet == FALSE || dwSize != ReadSize)
 675         {
 676                 Log("ReadFile error bRet:%u WriteSize:%u dwSize:%u ErrCode:%u\n", bRet, ReadSize, dwSize, GetLastError());
 677         }
 678
 679         return 1;
 680 }
 681
 682 static CHAR GetMountLogicalDrive(void)
 683 {
 684         CHAR Letter = 'Y';
 685         DWORD Drives;
 686         DWORD Mask = 0x1000000;
 687
 688         Drives = GetLogicalDrives();
 689     Log("Drives=0x%x", Drives);
 690
 691         while (Mask)
 692         {
 693                 if ((Drives & Mask) == 0)
 694                 {
 695                         break;
 696                 }
 697
 698                 Letter--;
 699                 Mask >>= 1;
 700         }
 701
 702         return Letter;
 703 }
 704
 705 UINT64 GetVentoyEfiPartStartSector(HANDLE hDrive)
 706 {
 707         BOOL bRet;
 708         DWORD dwSize;
 709         MBR_HEAD MBR;
 710         VTOY_GPT_INFO *pGpt = NULL;
 711         UINT64 StartSector = 0;
 712
 713         SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
 714
 715         bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
 716         Log("Read MBR Ret:%u Size:%u code:%u", bRet, dwSize, LASTERR);
 717
 718         if ((!bRet) || (dwSize != sizeof(MBR)))
 719         {
 720                 0;
 721         }
 722
 723         if (MBR.PartTbl[0].FsFlag == 0xEE)
 724         {
 725                 Log("GPT partition style");
 726
 727                 pGpt = malloc(sizeof(VTOY_GPT_INFO));
 728                 if (!pGpt)
 729                 {
 730                         return 0;
 731                 }
 732
 733                 SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
 734                 bRet = ReadFile(hDrive, pGpt, sizeof(VTOY_GPT_INFO), &dwSize, NULL);
 735                 if ((!bRet) || (dwSize != sizeof(VTOY_GPT_INFO)))
 736                 {
 737                         Log("Failed to read gpt info %d %u %d", bRet, dwSize, LASTERR);
 738                         return 0;
 739                 }
 740
 741                 StartSector = pGpt->PartTbl[1].StartLBA;
 742                 free(pGpt);
 743         }
 744         else
 745         {
 746                 Log("MBR partition style");
 747                 StartSector = MBR.PartTbl[1].StartSectorId;
 748         }
 749
 750         Log("GetVentoyEfiPart StartSector: %llu", StartSector);
 751         return StartSector;
 752 }
 753
 754 int VentoyMountISOByImdisk(const char *IsoPath, DWORD PhyDrive)
 755 {
 756         int rc = 1;
 757         BOOL bRet;
 758         CHAR Letter;
 759         DWORD dwBytes;
 760         HANDLE hDrive;
 761         CHAR PhyPath[MAX_PATH];
 762         WCHAR PhyPathW[MAX_PATH];
 763         STARTUPINFOA Si;
 764         PROCESS_INFORMATION Pi;
 765         GET_LENGTH_INFORMATION LengthInfo;
 766
 767         Log("VentoyMountISOByImdisk %s", IsoPath);
 768
 769         sprintf_s(PhyPath, sizeof(PhyPath), "\\\\.\\PhysicalDrive%d", PhyDrive);
 770     if (IsUTF8Encode(PhyPath))
 771     {
 772         Utf8ToUtf16(PhyPath, PhyPathW);
 773         hDrive = CreateFileW(PhyPathW, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 774     }
 775     else
 776     {
 777         hDrive = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 778     }
 779
 780         if (hDrive == INVALID_HANDLE_VALUE)
 781         {
 782                 Log("Could not open the disk<%s>, error:%u", PhyPath, GetLastError());
 783                 goto End;
 784         }
 785
 786         bRet = DeviceIoControl(hDrive, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &LengthInfo, sizeof(LengthInfo), &dwBytes, NULL);
 787         if (!bRet)
 788         {
 789                 Log("Could not get phy disk %s size, error:%u", PhyPath, GetLastError());
 790                 goto End;
 791         }
 792
 793         g_FatPhyDrive = hDrive;
 794         g_Part2StartSec = GetVentoyEfiPartStartSector(hDrive);
 795
 796         Log("Parse FAT fs...");
 797
 798         fl_init();
 799
 800         if (0 == fl_attach_media(VentoyFatDiskRead, NULL))
 801         {
 802                 if (g_64bit_system)
 803                 {
 804                         CopyFileFromFatDisk("/ventoy/imdisk/64/imdisk.sys", "ventoy\\imdisk.sys");
 805                         CopyFileFromFatDisk("/ventoy/imdisk/64/imdisk.exe", "ventoy\\imdisk.exe");
 806                         CopyFileFromFatDisk("/ventoy/imdisk/64/imdisk.cpl", "ventoy\\imdisk.cpl");
 807                 }
 808                 else
 809                 {
 810                         CopyFileFromFatDisk("/ventoy/imdisk/32/imdisk.sys", "ventoy\\imdisk.sys");
 811                         CopyFileFromFatDisk("/ventoy/imdisk/32/imdisk.exe", "ventoy\\imdisk.exe");
 812                         CopyFileFromFatDisk("/ventoy/imdisk/32/imdisk.cpl", "ventoy\\imdisk.cpl");
 813                 }
 814
 815                 GetCurrentDirectoryA(sizeof(PhyPath), PhyPath);
 816                 strcat_s(PhyPath, sizeof(PhyPath), "\\ventoy\\imdisk.sys");
 817
 818                 if (LoadNtDriver(PhyPath) == 0)
 819                 {
 820                         rc = 0;
 821
 822                         Letter = GetMountLogicalDrive();
 823             sprintf_s(PhyPath, sizeof(PhyPath), "ventoy\\imdisk.exe -a -o ro -f %s -m %C:", IsoPath, Letter);
 824
 825                         Log("mount iso to %C: use imdisk cmd <%s>", Letter, PhyPath);
 826
 827             GetStartupInfoA(&Si);
 828
 829             Si.dwFlags |= STARTF_USESHOWWINDOW;
 830             Si.wShowWindow = SW_HIDE;
 831
 832                         CreateProcessA(NULL, PhyPath, NULL, NULL, FALSE, 0, NULL, NULL, &Si, &Pi);
 833                         WaitForSingleObject(Pi.hProcess, INFINITE);
 834                 }
 835         }
 836         fl_shutdown();
 837
 838 End:
 839
 840         SAFE_CLOSE_HANDLE(hDrive);
 841
 842         return rc;
 843 }
 844
 845 static int MountIsoFile(CONST CHAR *IsoPath, DWORD PhyDrive)
 846 {
 847     if (IsWindows8OrGreater())
 848     {
 849         Log("This is Windows 8 or latter...");
 850         if (VentoyMountISOByAPI(IsoPath) == 0)
 851         {
 852             Log("Mount iso by API success");
 853             return 0;
 854         }
 855         else
 856         {
 857             Log("Mount iso by API failed, maybe not supported, try imdisk");
 858             return VentoyMountISOByImdisk(IsoPath, PhyDrive);
 859         }
 860     }
 861     else
 862     {
 863         Log("This is before Windows 8 ...");
 864         if (VentoyMountISOByImdisk(IsoPath, PhyDrive) == 0)
 865         {
 866             Log("Mount iso by imdisk success");
 867             return 0;
 868         }
 869         else
 870         {
 871             return VentoyMountISOByAPI(IsoPath);
 872         }
 873     }
 874 }
 875
 876 static int GetPhyDriveByLogicalDrive(int DriveLetter)
 877 {
 878     BOOL Ret;
 879     DWORD dwSize;
 880     HANDLE Handle;
 881     VOLUME_DISK_EXTENTS DiskExtents;
 882     CHAR PhyPath[128];
 883
 884     sprintf_s(PhyPath, sizeof(PhyPath), "\\\\.\\%C:", (CHAR)DriveLetter);
 885
 886     Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
 887     if (Handle == INVALID_HANDLE_VALUE)
 888     {
 889         Log("Could not open the disk<%s>, error:%u", PhyPath, GetLastError());
 890         return -1;
 891     }
 892
 893     Ret = DeviceIoControl(Handle,
 894         IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS,
 895         NULL,
 896         0,
 897         &DiskExtents,
 898         (DWORD)(sizeof(DiskExtents)),
 899         (LPDWORD)&dwSize,
 900         NULL);
 901
 902     if (!Ret || DiskExtents.NumberOfDiskExtents == 0)
 903     {
 904         Log("DeviceIoControl IOCTL_VOLUME_GET_VOLUME_DISK_EXTENTS failed %s, error:%u", PhyPath, GetLastError());
 905         SAFE_CLOSE_HANDLE(Handle);
 906         return -1;
 907     }
 908     SAFE_CLOSE_HANDLE(Handle);
 909
 910     Log("LogicalDrive:%s PhyDrive:%d Offset:%llu ExtentLength:%llu",
 911         PhyPath,
 912         DiskExtents.Extents[0].DiskNumber,
 913         DiskExtents.Extents[0].StartingOffset.QuadPart,
 914         DiskExtents.Extents[0].ExtentLength.QuadPart
 915         );
 916
 917     return (int)DiskExtents.Extents[0].DiskNumber;
 918 }
 919
 920
 921 static int DeleteVentoyPart2MountPoint(DWORD PhyDrive)
 922 {
 923     CHAR Letter = 'A';
 924     DWORD Drives;
 925     DWORD PhyDisk;
 926     CHAR DriveName[] = "?:\\";
 927
 928     Log("DeleteVentoyPart2MountPoint Phy%u ...", PhyDrive);
 929
 930     Drives = GetLogicalDrives();
 931     while (Drives)
 932     {
 933         if ((Drives & 0x01) && IsFileExist("%C:\\ventoy\\ventoy.cpio", Letter))
 934         {
 935             Log("File %C:\\ventoy\\ventoy.cpio exist", Letter);
 936
 937             PhyDisk = GetPhyDriveByLogicalDrive(Letter);
 938             Log("PhyDisk=%u for %C", PhyDisk, Letter);
 939
 940             if (PhyDisk == PhyDrive)
 941             {
 942                 DriveName[0] = Letter;
 943                 DeleteVolumeMountPointA(DriveName);
 944                 return 0;
 945             }
 946         }
 947
 948         Letter++;
 949         Drives >>= 1;
 950     }
 951
 952     return 1;
 953 }
 954
 955 static BOOL check_tar_archive(const char *archive, CHAR *tarName)
 956 {
 957     int len;
 958     int nameLen;
 959     const char *pos = archive;
 960     const char *slash = archive;
 961
 962     while (*pos)
 963     {
 964         if (*pos == '\\' || *pos == '/')
 965         {
 966             slash = pos;
 967         }
 968         pos++;
 969     }
 970
 971     len = (int)strlen(slash);
 972
 973     if (len > 7 && (strncmp(slash + len - 7, ".tar.gz", 7) == 0 || strncmp(slash + len - 7, ".tar.xz", 7) == 0))
 974     {
 975         nameLen = (int)sprintf_s(tarName, MAX_PATH, "X:%s", slash);
 976         tarName[nameLen - 3] = 0;
 977         return TRUE;
 978     }
 979     else if (len > 8 && strncmp(slash + len - 8, ".tar.bz2", 8) == 0)
 980     {
 981         nameLen = (int)sprintf_s(tarName, MAX_PATH, "X:%s", slash);
 982         tarName[nameLen - 4] = 0;
 983         return TRUE;
 984     }
 985     else if (len > 9 && strncmp(slash + len - 9, ".tar.lzma", 9) == 0)
 986     {
 987         nameLen = (int)sprintf_s(tarName, MAX_PATH, "X:%s", slash);
 988         tarName[nameLen - 5] = 0;
 989         return TRUE;
 990     }
 991
 992     return FALSE;
 993 }
 994
 995 static int DecompressInjectionArchive(const char *archive, DWORD PhyDrive)
 996 {
 997     int rc = 1;
 998     BOOL bRet;
 999     DWORD dwBytes;
1000     HANDLE hDrive;
1001     HANDLE hOut;
1002     DWORD flags = CREATE_NO_WINDOW;
1003     CHAR StrBuf[MAX_PATH];
1004     CHAR tarName[MAX_PATH];
1005     STARTUPINFOA Si;
1006     PROCESS_INFORMATION Pi;
1007     PROCESS_INFORMATION NewPi;
1008     GET_LENGTH_INFORMATION LengthInfo;
1009     SECURITY_ATTRIBUTES Sa = { sizeof(SECURITY_ATTRIBUTES), NULL, TRUE };
1010
1011     Log("DecompressInjectionArchive %s", archive);
1012
1013     sprintf_s(StrBuf, sizeof(StrBuf), "\\\\.\\PhysicalDrive%d", PhyDrive);
1014     hDrive = CreateFileA(StrBuf, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);
1015     if (hDrive == INVALID_HANDLE_VALUE)
1016     {
1017         Log("Could not open the disk<%s>, error:%u", StrBuf, GetLastError());
1018         goto End;
1019     }
1020
1021     bRet = DeviceIoControl(hDrive, IOCTL_DISK_GET_LENGTH_INFO, NULL, 0, &LengthInfo, sizeof(LengthInfo), &dwBytes, NULL);
1022     if (!bRet)
1023     {
1024         Log("Could not get phy disk %s size, error:%u", StrBuf, GetLastError());
1025         goto End;
1026     }
1027
1028     g_FatPhyDrive = hDrive;
1029         g_Part2StartSec = GetVentoyEfiPartStartSector(hDrive);
1030
1031     Log("Parse FAT fs...");
1032
1033     fl_init();
1034
1035     if (0 == fl_attach_media(VentoyFatDiskRead, NULL))
1036     {
1037         if (g_64bit_system)
1038         {
1039             CopyFileFromFatDisk("/ventoy/7z/64/7za.exe", "ventoy\\7za.exe");
1040         }
1041         else
1042         {
1043             CopyFileFromFatDisk("/ventoy/7z/32/7za.exe", "ventoy\\7za.exe");
1044         }
1045
1046         sprintf_s(StrBuf, sizeof(StrBuf), "ventoy\\7za.exe x -y -aoa -oX:\\ %s", archive);
1047
1048         Log("extract inject to X:");
1049         Log("cmdline:<%s>", StrBuf);
1050
1051         GetStartupInfoA(&Si);
1052
1053         hOut = CreateFileA("ventoy\\7z.log",
1054             FILE_APPEND_DATA,
1055             FILE_SHARE_WRITE | FILE_SHARE_READ,
1056             &Sa,
1057             OPEN_ALWAYS,
1058             FILE_ATTRIBUTE_NORMAL,
1059             NULL);
1060
1061         Si.dwFlags |= STARTF_USESTDHANDLES;
1062
1063         if (hOut != INVALID_HANDLE_VALUE)
1064         {
1065             Si.hStdError = hOut;
1066             Si.hStdOutput = hOut;
1067         }
1068
1069         CreateProcessA(NULL, StrBuf, NULL, NULL, TRUE, flags, NULL, NULL, &Si, &Pi);
1070         WaitForSingleObject(Pi.hProcess, INFINITE);
1071
1072         //
1073         // decompress tar archive, for tar.gz/tar.xz/tar.bz2
1074         //
1075         if (check_tar_archive(archive, tarName))
1076         {
1077             Log("Decompress tar archive...<%s>", tarName);
1078
1079             sprintf_s(StrBuf, sizeof(StrBuf), "ventoy\\7za.exe x -y -aoa -oX:\\ %s", tarName);
1080
1081             CreateProcessA(NULL, StrBuf, NULL, NULL, TRUE, flags, NULL, NULL, &Si, &NewPi);
1082             WaitForSingleObject(NewPi.hProcess, INFINITE);
1083
1084             Log("Now delete %s", tarName);
1085             DeleteFileA(tarName);
1086         }
1087
1088         SAFE_CLOSE_HANDLE(hOut);
1089     }
1090     fl_shutdown();
1091
1092 End:
1093
1094     SAFE_CLOSE_HANDLE(hDrive);
1095
1096     return rc;
1097 }
1098
1099 static int ProcessUnattendedInstallation(const char *script)
1100 {
1101     DWORD dw;
1102     HKEY hKey;
1103     LSTATUS Ret;
1104     CHAR Letter;
1105     CHAR CurDir[MAX_PATH];
1106
1107     Log("Copy unattended XML ...");
1108
1109     GetCurrentDirectory(sizeof(CurDir), CurDir);
1110     Letter = CurDir[0];
1111     if ((Letter >= 'A' && Letter <= 'Z') || (Letter >= 'a' && Letter <= 'z'))
1112     {
1113         Log("Current Drive Letter: %C", Letter);
1114     }
1115     else
1116     {
1117         Letter = 'X';
1118     }
1119
1120     sprintf_s(CurDir, sizeof(CurDir), "%C:\\Autounattend.xml", Letter);
1121     Log("Copy file <%s> --> <%s>", script, CurDir);
1122     CopyFile(script, CurDir, FALSE);
1123
1124     Ret = RegCreateKeyEx(HKEY_LOCAL_MACHINE, "System\\Setup", 0, NULL, REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, NULL, &hKey, &dw);
1125     if (ERROR_SUCCESS == Ret)
1126     {
1127         Ret = RegSetValueEx(hKey, "UnattendFile", 0, REG_SZ, CurDir, (DWORD)(strlen(CurDir) + 1));
1128     }
1129
1130     return 0;
1131 }
1132
1133 static int VentoyHook(ventoy_os_param *param)
1134 {
1135     int rc;
1136         CHAR Letter = 'A';
1137         DISK_EXTENT DiskExtent;
1138         DWORD Drives = GetLogicalDrives();
1139         UINT8 UUID[16];
1140         CHAR IsoPath[MAX_PATH];
1141
1142         Log("Logical Drives=0x%x Path:<%s>", Drives, param->vtoy_img_path);
1143
1144     if (IsUTF8Encode(param->vtoy_img_path))
1145     {
1146         Log("This file is UTF8 encoding\n");
1147     }
1148
1149         while (Drives)
1150         {
1151         if (Drives & 0x01)
1152         {
1153             sprintf_s(IsoPath, sizeof(IsoPath), "%C:\\%s", Letter, param->vtoy_img_path);
1154             if (IsFileExist("%s", IsoPath))
1155             {
1156                 Log("File exist under %C:", Letter);
1157                 if (GetPhyDiskUUID(Letter, UUID, &DiskExtent) == 0)
1158                 {
1159                     if (memcmp(UUID, param->vtoy_disk_guid, 16) == 0)
1160                     {
1161                         Log("Disk UUID match");
1162                         break;
1163                     }
1164                 }
1165             }
1166             else
1167             {
1168                 Log("File NOT exist under %C:", Letter);
1169             }
1170         }
1171
1172                 Drives >>= 1;
1173                 Letter++;
1174         }
1175
1176         if (Drives == 0)
1177         {
1178                 Log("Failed to find ISO file");
1179                 return 1;
1180         }
1181
1182         Log("Find ISO file <%s>", IsoPath);
1183
1184     rc = MountIsoFile(IsoPath, DiskExtent.DiskNumber);
1185     Log("Mount ISO FILE: %s", rc == 0 ? "SUCCESS" : "FAILED");
1186
1187     // for protect
1188     rc = DeleteVentoyPart2MountPoint(DiskExtent.DiskNumber);
1189     Log("Delete ventoy mountpoint: %s", rc == 0 ? "SUCCESS" : "NO NEED");
1190
1191     if (g_windows_data.auto_install_script[0])
1192     {
1193         sprintf_s(IsoPath, sizeof(IsoPath), "%C:%s", Letter, g_windows_data.auto_install_script);
1194         if (IsFileExist("%s", IsoPath))
1195         {
1196             Log("use auto install script %s...", IsoPath);
1197             ProcessUnattendedInstallation(IsoPath);
1198         }
1199         else
1200         {
1201             Log("auto install script %s not exist", IsoPath);
1202         }
1203     }
1204     else
1205     {
1206         Log("auto install no need");
1207     }
1208
1209     if (g_windows_data.injection_archive[0])
1210     {
1211         sprintf_s(IsoPath, sizeof(IsoPath), "%C:%s", Letter, g_windows_data.injection_archive);
1212         if (IsFileExist("%s", IsoPath))
1213         {
1214             Log("decompress injection archive %s...", IsoPath);
1215             DecompressInjectionArchive(IsoPath, DiskExtent.DiskNumber);
1216         }
1217         else
1218         {
1219             Log("injection archive %s not exist", IsoPath);
1220         }
1221     }
1222     else
1223     {
1224         Log("no injection archive found");
1225     }
1226
1227     return 0;
1228 }
1229
1230 const char * GetFileNameInPath(const char *fullpath)
1231 {
1232         int i;
1233         const char *pos = NULL;
1234
1235         if (strstr(fullpath, ":"))
1236         {
1237                 for (i = (int)strlen(fullpath); i > 0; i--)
1238                 {
1239                         if (fullpath[i - 1] == '/' || fullpath[i - 1] == '\\')
1240                         {
1241                                 return fullpath + i;
1242                         }
1243                 }
1244         }
1245
1246         return fullpath;
1247 }
1248
1249 int VentoyJumpWimboot(INT argc, CHAR **argv, CHAR *LunchFile)
1250 {
1251     int rc = 1;
1252     char *buf = NULL;
1253     DWORD size = 0;
1254     DWORD Pos;
1255
1256 #ifdef VTOY_32
1257     g_64bit_system = FALSE;
1258 #else
1259     g_64bit_system = TRUE;
1260 #endif
1261
1262     Log("VentoyJumpWimboot %dbit", g_64bit_system ? 64 : 32);
1263
1264     sprintf_s(LunchFile, MAX_PATH, "X:\\setup.exe");
1265
1266     ReadWholeFile2Buf("wimboot.data", &buf, &size);
1267     Log("wimboot.data size:%d", size);
1268
1269     memcpy(&g_os_param, buf, sizeof(ventoy_os_param));
1270     memcpy(&g_windows_data, buf + sizeof(ventoy_os_param), sizeof(ventoy_windows_data));
1271     memcpy(g_os_param_reserved, g_os_param.vtoy_reserved, sizeof(g_os_param_reserved));
1272
1273     if (g_os_param_reserved[0] == 1)
1274     {
1275         Log("break here for debug .....");
1276         goto End;
1277     }
1278
1279     // convert / to \\
1280     for (Pos = 0; Pos < sizeof(g_os_param.vtoy_img_path) && g_os_param.vtoy_img_path[Pos]; Pos++)
1281     {
1282         if (g_os_param.vtoy_img_path[Pos] == '/')
1283         {
1284             g_os_param.vtoy_img_path[Pos] = '\\';
1285         }
1286     }
1287
1288     if (g_os_param_reserved[0] == 2)
1289     {
1290         Log("skip hook for debug .....");
1291         rc = 0;
1292         goto End;
1293     }
1294
1295     rc = VentoyHook(&g_os_param);
1296
1297 End:
1298
1299     if (buf)
1300     {
1301         free(buf);
1302     }
1303
1304     return rc;
1305 }
1306
1307 int VentoyJump(INT argc, CHAR **argv, CHAR *LunchFile)
1308 {
1309         int rc = 1;
1310         DWORD Pos;
1311         DWORD PeStart;
1312     DWORD FileSize;
1313         BYTE *Buffer = NULL;
1314         CHAR ExeFileName[MAX_PATH];
1315
1316         sprintf_s(ExeFileName, sizeof(ExeFileName), "%s", argv[0]);
1317         if (!IsFileExist("%s", ExeFileName))
1318         {
1319                 Log("File %s NOT exist, now try %s.exe", ExeFileName, ExeFileName);
1320                 sprintf_s(ExeFileName, sizeof(ExeFileName), "%s.exe", argv[0]);
1321
1322                 Log("File %s exist ? %s", ExeFileName, IsFileExist("%s", ExeFileName) ? "YES" : "NO");
1323         }
1324
1325         if (ReadWholeFile2Buf(ExeFileName, (void **)&Buffer, &FileSize))
1326         {
1327                 goto End;
1328         }
1329
1330         g_64bit_system = IsPe64(Buffer);
1331     Log("VentoyJump %dbit", g_64bit_system ? 64 : 32);
1332
1333     if (IsDirExist("ventoy"))
1334     {
1335         Log("ventoy directory already exist");
1336     }
1337     else
1338         {
1339         Log("ventoy directory not exist, now create it.");
1340                 if (!CreateDirectoryA("ventoy", NULL))
1341                 {
1342                         Log("Failed to create ventoy directory err:%u", GetLastError());
1343                         goto End;
1344                 }
1345         }
1346
1347         for (PeStart = 0; PeStart < FileSize; PeStart += 16)
1348         {
1349                 if (CheckOsParam((ventoy_os_param *)(Buffer + PeStart)) &&
1350             CheckPeHead(Buffer + PeStart + sizeof(ventoy_os_param) + sizeof(ventoy_windows_data)))
1351                 {
1352                         Log("Find os pararm at %u", PeStart);
1353
1354             memcpy(&g_os_param, Buffer + PeStart, sizeof(ventoy_os_param));
1355             memcpy(&g_windows_data, Buffer + PeStart + sizeof(ventoy_os_param), sizeof(ventoy_windows_data));
1356             memcpy(g_os_param_reserved, g_os_param.vtoy_reserved, sizeof(g_os_param_reserved));
1357
1358             if (g_os_param_reserved[0] == 1)
1359                         {
1360                                 Log("break here for debug .....");
1361                                 goto End;
1362                         }
1363
1364                         // convert / to \\
1365             for (Pos = 0; Pos < sizeof(g_os_param.vtoy_img_path) && g_os_param.vtoy_img_path[Pos]; Pos++)
1366                         {
1367                 if (g_os_param.vtoy_img_path[Pos] == '/')
1368                                 {
1369                     g_os_param.vtoy_img_path[Pos] = '\\';
1370                                 }
1371                         }
1372
1373                         PeStart += sizeof(ventoy_os_param) + sizeof(ventoy_windows_data);
1374                         sprintf_s(LunchFile, MAX_PATH, "ventoy\\%s", GetFileNameInPath(ExeFileName));
1375
1376             if (IsFileExist("%s", LunchFile))
1377             {
1378                 Log("vtoyjump multiple call...");
1379                 rc = 0;
1380                 goto End;
1381             }
1382
1383                         SaveBuffer2File(LunchFile, Buffer + PeStart, FileSize - PeStart);
1384                         break;
1385                 }
1386         }
1387
1388         if (PeStart >= FileSize)
1389         {
1390                 Log("OS param not found");
1391                 goto End;
1392         }
1393
1394     if (g_os_param_reserved[0] == 2)
1395     {
1396         Log("skip hook for debug .....");
1397         rc = 0;
1398         goto End;
1399     }
1400
1401     rc = VentoyHook(&g_os_param);
1402
1403 End:
1404
1405         if (Buffer)
1406         {
1407                 free(Buffer);
1408         }
1409
1410         return rc;
1411 }
1412
1413 static int GetPecmdParam(const char *argv, char *CallParamBuf, DWORD BufLen)
1414 {
1415     HKEY hKey;
1416     LSTATUS Ret;
1417     DWORD dw;
1418     DWORD Type;
1419     CHAR *Pos = NULL;
1420     CHAR CallParam[256] = { 0 };
1421     CHAR FileName[MAX_PATH];
1422
1423     Log("GetPecmdParam <%s>", argv);
1424
1425     *CallParamBuf = 0;
1426
1427     strcpy_s(FileName, sizeof(FileName), argv);
1428     for (dw = 0, Pos = FileName; *Pos; Pos++)
1429     {
1430         dw++;
1431         *Pos = toupper(*Pos);
1432     }
1433
1434     Log("dw=%lu argv=<%s>", dw, FileName);
1435
1436     if (dw >= 9 && strcmp(FileName + dw - 9, "PECMD.EXE") == 0)
1437     {
1438         Log("Get parameters for pecmd.exe");
1439         Ret = RegCreateKeyEx(HKEY_LOCAL_MACHINE, "System\\Setup", 0, NULL, REG_OPTION_NON_VOLATILE, KEY_ALL_ACCESS, NULL, &hKey, &dw);
1440         if (ERROR_SUCCESS == Ret)
1441         {
1442             memset(FileName, 0, sizeof(FileName));
1443             dw = sizeof(FileName);
1444             Ret = RegQueryValueEx(hKey, "CmdLine", NULL, &Type, FileName, &dw);
1445             if (ERROR_SUCCESS == Ret && Type == REG_SZ)
1446             {
1447                 strcpy_s(CallParam, sizeof(CallParam), FileName);
1448                 Log("CmdLine:<%s>", CallParam);
1449
1450                 if (_strnicmp(CallParam, "PECMD.EXE", 9) == 0)
1451                 {
1452                     Pos = CallParam + 9;
1453                     if (*Pos == ' ' || *Pos == '\t')
1454                     {
1455                         Pos++;
1456                     }
1457                 }
1458                 else
1459                 {
1460                     Pos = CallParam;
1461                 }
1462
1463                 Log("CmdLine2:<%s>", Pos);
1464                 sprintf_s(CallParamBuf, BufLen, " %s", Pos);
1465             }
1466             else
1467             {
1468                 Log("Failed to RegQueryValueEx %lu %lu", Ret, Type);
1469             }
1470
1471             RegCloseKey(hKey);
1472             return 1;
1473         }
1474         else
1475         {
1476             Log("Failed to create reg key %lu", Ret);
1477         }
1478     }
1479     else
1480     {
1481         Log("This is NOT pecmd.exe");
1482     }
1483
1484     return 0;
1485 }
1486
1487 static int GetWpeInitParam(char **argv, int argc, char *CallParamBuf, DWORD BufLen)
1488 {
1489     int i;
1490     DWORD dw;
1491     CHAR *Pos = NULL;
1492     CHAR FileName[MAX_PATH];
1493
1494     Log("GetWpeInitParam argc=%d", argc);
1495
1496     *CallParamBuf = 0;
1497
1498     strcpy_s(FileName, sizeof(FileName), argv[0]);
1499     for (dw = 0, Pos = FileName; *Pos; Pos++)
1500     {
1501         dw++;
1502         *Pos = toupper(*Pos);
1503     }
1504
1505     Log("dw=%lu argv=<%s>", dw, FileName);
1506
1507     if (dw >= 11 && strcmp(FileName + dw - 11, "WPEINIT.EXE") == 0)
1508     {
1509         Log("Get parameters for WPEINIT.EXE");
1510         for (i = 1; i < argc; i++)
1511         {
1512             strcat_s(CallParamBuf, BufLen, " ");
1513             strcat_s(CallParamBuf, BufLen, argv[i]);
1514         }
1515
1516         return 1;
1517     }
1518     else
1519     {
1520         Log("This is NOT wpeinit.exe");
1521     }
1522
1523     return 0;
1524 }
1525
1526
1527 int main(int argc, char **argv)
1528 {
1529     int i = 0;
1530     int rc = 0;
1531         CHAR *Pos = NULL;
1532         CHAR CurDir[MAX_PATH];
1533         CHAR LunchFile[MAX_PATH];
1534     CHAR CallParam[1024] = { 0 };
1535         STARTUPINFOA Si;
1536         PROCESS_INFORMATION Pi;
1537
1538         if (argv[0] && argv[0][0] && argv[0][1] == ':')
1539         {
1540                 GetCurrentDirectoryA(sizeof(CurDir), CurDir);
1541
1542                 strcpy_s(LunchFile, sizeof(LunchFile), argv[0]);
1543                 Pos = (char *)GetFileNameInPath(LunchFile);
1544
1545                 strcat_s(CurDir, sizeof(CurDir), "\\");
1546                 strcat_s(CurDir, sizeof(CurDir), Pos);
1547
1548                 if (_stricmp(argv[0], CurDir) != 0)
1549                 {
1550                         *Pos = 0;
1551                         SetCurrentDirectoryA(LunchFile);
1552                 }
1553         }
1554
1555         Log("######## VentoyJump ##########");
1556         Log("argc = %d argv[0] = <%s>", argc, argv[0]);
1557
1558         if (Pos && *Pos == 0)
1559         {
1560                 Log("Old current directory = <%s>", CurDir);
1561                 Log("New current directory = <%s>", LunchFile);
1562         }
1563         else
1564         {
1565                 GetCurrentDirectoryA(sizeof(CurDir), CurDir);
1566                 Log("Current directory = <%s>", CurDir);
1567         }
1568
1569     if (0 == GetWpeInitParam(argv, argc, CallParam, sizeof(CallParam)))
1570     {
1571         GetPecmdParam(argv[0], CallParam, sizeof(CallParam));
1572     }
1573
1574     GetStartupInfoA(&Si);
1575
1576     memset(LunchFile, 0, sizeof(LunchFile));
1577
1578     if (strstr(argv[0], "vtoyjump.exe"))
1579     {
1580         rc = VentoyJumpWimboot(argc, argv, LunchFile);
1581     }
1582     else
1583     {
1584         rc = VentoyJump(argc, argv, LunchFile);
1585     }
1586
1587     Log("LunchFile=<%s> CallParam=<%s>", LunchFile, CallParam);
1588
1589     if (g_os_param_reserved[0] == 3)
1590     {
1591         Log("Open log for debug ...");
1592         sprintf_s(LunchFile, sizeof(LunchFile), "%s", "notepad.exe ventoy.log");
1593     }
1594     else
1595     {
1596         if (CallParam[0])
1597         {
1598             strcat_s(LunchFile, sizeof(LunchFile), CallParam);
1599         }
1600         else if (NULL == strstr(LunchFile, "setup.exe"))
1601         {
1602             Log("Not setup.exe, hide windows.");
1603             Si.dwFlags |= STARTF_USESHOWWINDOW;
1604             Si.wShowWindow = SW_HIDE;
1605         }
1606
1607         Log("Ventoy jump %s ...", rc == 0 ? "success" : "failed");
1608     }
1609
1610     Log("Now launch <%s> ...", LunchFile);
1611
1612     //sprintf_s(LunchFile, sizeof(LunchFile), "%s", "cmd.exe");
1613     CreateProcessA(NULL, LunchFile, NULL, NULL, FALSE, 0, NULL, NULL, &Si, &Pi);
1614
1615     for (i = 0; rc && i < 10; i++)
1616     {
1617         Log("Ventoy hook failed, now wait and retry ...");
1618         Sleep(1000);
1619         rc = VentoyHook(&g_os_param);
1620     }
1621
1622     Log("Wait process...");
1623     WaitForSingleObject(Pi.hProcess, INFINITE);
1624
1625     Log("vtoyjump finished");
1626         return 0;
1627 }