vtoyjump/vtoyjump/setupmon.c

   1 /******************************************************************************
   2 * setupmon.c
   3 *
   4 * Copyright (c) 2022, 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
  30
  31 #define PROMPT_MAX  1024
  32 typedef struct VTOY_MUI_DATA
  33 {
  34     WCHAR Prompt[PROMPT_MAX];
  35     struct VTOY_MUI_DATA* pNext;
  36 }VTOY_MUI_DATA;
  37
  38 static VTOY_MUI_DATA* g_MuiDataHead = NULL;
  39 static FILETIME g_SetupMonStartTime;
  40
  41 #pragma pack(1)
  42 typedef struct {
  43     WORD      dlgVer;
  44     WORD      signature;
  45     DWORD     helpID;
  46     DWORD     exStyle;
  47     DWORD     style;
  48     WORD      cDlgItems;
  49     short     x;
  50     short     y;
  51     short     cx;
  52     short     cy;
  53     //sz_Or_Ord menu;
  54     //sz_Or_Ord windowClass;
  55     //WCHAR     title[titleLen];
  56     //WORD      pointsize;
  57     //WORD      weight;
  58     //BYTE      italic;
  59     //BYTE      charset;
  60     //WCHAR     typeface[stringLen];
  61 } DLGTEMPLATEEX;
  62
  63 typedef struct {
  64     DWORD     helpID;
  65     DWORD     exStyle;
  66     DWORD     style;
  67     short     x;
  68     short     y;
  69     short     cx;
  70     short     cy;
  71     DWORD     id;
  72     //  sz_Or_Ord windowClass;
  73     //  sz_Or_Ord title;
  74     //  WORD      extraCount;
  75 } DLGITEMTEMPLATEEX;
  76
  77
  78 #pragma pack()
  79
  80 #define wchar_to_utf8_no_alloc(wsrc, dest, dest_size) \
  81         WideCharToMultiByte(CP_UTF8, 0, wsrc, -1, dest, dest_size, NULL, NULL)
  82 #define utf8_to_wchar_no_alloc(src, wdest, wdest_size) \
  83         MultiByteToWideChar(CP_UTF8, 0, src, -1, wdest, wdest_size)
  84
  85 #define VTOY_DLG_SKIP_SZ_ORD(pWord) \
  86 {\
  87     if (*pWord == 0x0000)\
  88     {\
  89         pWordData += 1;\
  90     }\
  91     else if (*pWordData == 0xFFFF)\
  92     {\
  93         pWordData += 2;\
  94     }\
  95     else\
  96     {\
  97         while (*pWordData++)\
  98         {\
  99             ;\
 100         }\
 101     }\
 102 }
 103
 104 #define VTOY_DLG_SKIP_SZ(pWord) \
 105 {\
 106     while (*pWord)\
 107     {\
 108         pWord++;\
 109     }\
 110     pWord++;\
 111 }
 112
 113
 114 static BOOL LoadSetupRebootDialogPrompt(HMODULE hMui, DWORD ID, WCHAR* Prompt, DWORD Len)
 115 {
 116     DWORD i = 0;
 117     UINT64 Addr = 0;
 118     WORD* pWordData = NULL;
 119     HRSRC hDlg = NULL;
 120     HGLOBAL hTemplate = NULL;
 121     DLGTEMPLATEEX* pDlgTempEx = NULL;
 122     DLGITEMTEMPLATEEX* pDlgItemTempEx = NULL;
 123
 124     hDlg = FindResource(hMui, MAKEINTRESOURCE(ID), RT_DIALOG);
 125     if (!hDlg)
 126     {
 127         return FALSE;
 128     }
 129
 130     hTemplate = LoadResource(hMui, hDlg);
 131     if (!hTemplate)
 132     {
 133         return FALSE;
 134     }
 135
 136     pDlgTempEx = (DLGTEMPLATEEX*)LockResource(hTemplate);
 137     if (!pDlgTempEx)
 138     {
 139         return FALSE;
 140     }
 141
 142     if (pDlgTempEx->signature != 0xFFFF)
 143     {
 144         return FALSE;
 145     }
 146
 147     pWordData = (WORD*)(pDlgTempEx + 1);
 148
 149     //skip menu
 150     VTOY_DLG_SKIP_SZ_ORD(pWordData);
 151
 152     //skip windowClass
 153     VTOY_DLG_SKIP_SZ_ORD(pWordData);
 154
 155     //skip title
 156     VTOY_DLG_SKIP_SZ(pWordData);
 157
 158     //skip pointsize/weight/italic + charset
 159     pWordData += 3;
 160
 161     //skip typeface
 162     VTOY_DLG_SKIP_SZ(pWordData);
 163
 164     //align to DWORD
 165     Addr = (UINT64)pWordData;
 166     if ((Addr % 4) > 0)
 167     {
 168         Addr += 4 - (Addr % 4);
 169         pWordData = (WORD*)Addr;
 170     }
 171
 172     //First Dlg Item
 173     pDlgItemTempEx = (DLGITEMTEMPLATEEX*)pWordData;
 174
 175     if (pDlgItemTempEx->id != 1023)
 176     {
 177         return FALSE;
 178     }
 179
 180     pWordData = (WORD*)(pDlgItemTempEx + 1);
 181
 182     //skip windowClass
 183     VTOY_DLG_SKIP_SZ_ORD(pWordData);
 184
 185     for (i = 0; i < Len; i++)
 186     {
 187         Prompt[i] = pWordData[i];
 188         if (Prompt[i] == 0)
 189         {
 190             break;
 191         }
 192     }
 193
 194     return TRUE;
 195 }
 196
 197 static int LoadPromptString(const char* szDir)
 198 {
 199     BOOL bRet;
 200     HMODULE hMui;
 201     CHAR MuiFile[MAX_PATH];
 202     WCHAR Prompt[PROMPT_MAX];
 203     VTOY_MUI_DATA* Node = NULL;
 204
 205     sprintf_s(MuiFile, sizeof(MuiFile), "X:\\Sources\\%s\\w32uires.dll.mui", szDir);
 206
 207     hMui = LoadLibraryA(MuiFile);
 208     if (!hMui)
 209     {
 210         Log("Failed to loadlibrary <%s> %u", MuiFile, LASTERR);
 211         return 1;
 212     }
 213
 214     bRet = LoadSetupRebootDialogPrompt(hMui, 103, Prompt, PROMPT_MAX);
 215     if (bRet)
 216     {
 217         Node = malloc(sizeof(VTOY_MUI_DATA));
 218         if (Node)
 219         {
 220             memset(Node, 0, sizeof(VTOY_MUI_DATA));
 221             memcpy(Node->Prompt, Prompt, sizeof(Prompt));
 222             if (g_MuiDataHead == NULL)
 223             {
 224                 g_MuiDataHead = Node;
 225             }
 226             else
 227             {
 228                 Node->pNext = g_MuiDataHead;
 229                 g_MuiDataHead = Node;
 230             }
 231
 232             Log("Successfully add prompt string for <%s>", szDir);
 233         }
 234     }
 235     else
 236     {
 237         Log("Failed to load prompt string from %s", szDir);
 238     }
 239
 240     FreeLibrary(hMui);
 241     return 0;
 242 }
 243
 244 static int FindAllPromptStrings(void)
 245 {
 246     HANDLE hFind;
 247     WIN32_FIND_DATAA FindData;
 248
 249     hFind = FindFirstFileA("X:\\Sources\\*", &FindData);
 250     if (hFind == INVALID_HANDLE_VALUE)
 251     {
 252         Log("FindFirstFileA failed %u", LASTERR);
 253         return 0;
 254     }
 255
 256     do
 257     {
 258         if ((FindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) && (strlen(FindData.cFileName) == 5))
 259         {
 260             if (IsFileExist("X:\\Sources\\%s\\w32uires.dll.mui", FindData.cFileName))
 261             {
 262                 LoadPromptString(FindData.cFileName);
 263             }
 264         }
 265     } while (FindNextFileA(hFind, &FindData));
 266
 267     return 0;
 268 }
 269
 270 static BOOL CALLBACK SetupPromptCallback(HWND hWnd, LPARAM lParam)
 271 {
 272     VTOY_MUI_DATA* Node = NULL;
 273     WCHAR Prompt[PROMPT_MAX] = { 0 };
 274     BOOL* found = (BOOL*)lParam;
 275
 276     if (GetWindowTextW(hWnd, Prompt, PROMPT_MAX - 1) > 0)
 277     {
 278         for (Node = g_MuiDataHead; Node; Node = Node->pNext)
 279         {
 280             if (wcscmp(Node->Prompt, Prompt) == 0)
 281             {
 282                 *found = TRUE;
 283                 break;
 284             }
 285         }
 286     }
 287
 288     return TRUE;
 289 }
 290
 291 static BOOL CALLBACK EnumWindowsProcFunc(HWND hWnd, LPARAM lParam)
 292 {
 293     EnumChildWindows(hWnd, SetupPromptCallback, lParam);
 294     return TRUE;
 295 }
 296
 297
 298 static CHAR FindWindowsInstallDstDrive(void)
 299 {
 300     CHAR ret = 0;
 301     CHAR Drv = 'A';
 302     DWORD Drives;
 303     CHAR ChkFilePath[MAX_PATH];
 304     HANDLE hFile;
 305     UINT64 Time1, Time2, Time3;
 306     FILETIME fTime;
 307     FILETIME sysftime;
 308     SYSTEMTIME cur_systime;
 309
 310     GetSystemTime(&cur_systime);
 311     SystemTimeToFileTime(&cur_systime, &sysftime);
 312
 313     memcpy(&Time1, &g_SetupMonStartTime, 8);
 314     memcpy(&Time2, &sysftime, 8);
 315
 316     Drives = GetLogicalDrives();
 317     Log("Find Windows: Logical Drives: 0x%x", Drives);
 318
 319     while ((Drives > 0) && (ret == 0))
 320     {
 321         if ((Drives & 1) && Drv != 'X')
 322         {
 323             sprintf_s(ChkFilePath, sizeof(ChkFilePath), "%C:\\$WINDOWS.~BT\\Sources\\Panther\\setupact.log", Drv);
 324             hFile = CreateFileA(ChkFilePath, FILE_READ_EA, FILE_SHARE_READ, 0, OPEN_EXISTING, 0, 0);
 325             if (hFile != INVALID_HANDLE_VALUE)
 326             {
 327                 Log("%s Exist", ChkFilePath);
 328
 329                 if (GetFileTime(hFile, NULL, NULL, &fTime))
 330                 {
 331                     memcpy(&Time3, &fTime, 8);
 332
 333                     Log("### %s %llu %llu %llu", ChkFilePath, Time1, Time2, Time3);
 334
 335                     if (Time3 > Time1 && Time3 < Time2)
 336                     {
 337                         Log("Detect Windows11 install drive:<%C>", Drv);
 338                         ret = Drv;
 339                     }
 340                 }
 341                 CloseHandle(hFile);
 342             }
 343             else
 344             {
 345                 Log("%s NOT Exist %u", ChkFilePath, LASTERR);
 346             }
 347         }
 348
 349         Drives >>= 1;
 350         Drv++;
 351     }
 352
 353     return ret;
 354 }
 355
 356 static int AddBypassNROReg(const CHAR* OfflineRegPath)
 357 {
 358     int ret = 1;
 359     BOOL Unload = FALSE;
 360     LSTATUS Status;
 361     HKEY hKey = NULL;
 362     HKEY hSubKey = NULL;
 363     DWORD dwValue = 1;
 364     DWORD dwSize;
 365     HANDLE htoken;
 366     DWORD s = 0;
 367     TOKEN_PRIVILEGES* p = NULL;
 368
 369     Log("AddBypassNROReg<%s>", OfflineRegPath);
 370
 371     if (!IsFileExist(OfflineRegPath))
 372     {
 373         return ret;
 374     }
 375
 376     if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &htoken))
 377     {
 378         Log("Open process token failed! %u", LASTERR);
 379         return ret;
 380     }
 381
 382     s = sizeof(TOKEN_PRIVILEGES) + 2 * sizeof(LUID_AND_ATTRIBUTES);
 383     p = (PTOKEN_PRIVILEGES)malloc(s);
 384     if (!p)
 385     {
 386         Log("Failed to alloc privileges memory");
 387         return ret;
 388     }
 389
 390     if (!LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &(p->Privileges[0].Luid)) ||
 391         !LookupPrivilegeValue(NULL, SE_BACKUP_NAME, &(p->Privileges[1].Luid)) ||
 392         !LookupPrivilegeValue(NULL, SE_RESTORE_NAME, &(p->Privileges[2].Luid)))
 393     {
 394         Log("LookupPrivilegeValue failed!");
 395         goto End;
 396     }
 397
 398     p->PrivilegeCount = 3;
 399     p->Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
 400     p->Privileges[1].Attributes = SE_PRIVILEGE_ENABLED;
 401     p->Privileges[2].Attributes = SE_PRIVILEGE_ENABLED;
 402
 403     if (!AdjustTokenPrivileges(htoken, FALSE, p, s, NULL, NULL) || GetLastError() != ERROR_SUCCESS)
 404     {
 405         Log("AdjustTokenPrivileges failed!");
 406         goto End;
 407     }
 408
 409     Log("AdjustTokenPrivileges success");
 410
 411     Status = RegLoadKeyA(HKEY_LOCAL_MACHINE, "VTOYNEWSW", OfflineRegPath);
 412     if (Status != ERROR_SUCCESS)
 413     {
 414         Log("RegLoadKey Failed 0x%x", Status);
 415         goto End;
 416     }
 417     Unload = TRUE;
 418
 419     Status = RegCreateKeyExA(HKEY_LOCAL_MACHINE, "VTOYNEWSW\\Microsoft\\Windows\\CurrentVersion", 0, NULL, REG_OPTION_BACKUP_RESTORE, KEY_ALL_ACCESS, NULL, &hKey, &dwSize);
 420     if (ERROR_SUCCESS != Status)
 421     {
 422         Log("Failed to create reg key VTOYNEWSW\\Microsoft\\Windows\\CurrentVersion %u %u", LASTERR, Status);
 423         goto End;
 424     }
 425
 426     Status = RegCreateKeyExA(hKey, "OOBE", 0, NULL, 0, KEY_SET_VALUE | KEY_QUERY_VALUE | KEY_CREATE_SUB_KEY, NULL, &hSubKey, &dwSize);
 427     if (ERROR_SUCCESS != Status)
 428     {
 429         Log("Failed to create OOBE reg  %u %u", LASTERR, Status);
 430         goto End;
 431     }
 432
 433     Status = RegSetValueExA(hSubKey, "BypassNRO", 0, REG_DWORD, (LPBYTE)&dwValue, sizeof(DWORD));
 434     Log("Create BypassNRO registry %s %u", (Status == ERROR_SUCCESS) ? "SUCCESS" : "FAILED", Status);
 435
 436     Status = RegFlushKey(hSubKey);
 437     Status += RegCloseKey(hSubKey);
 438     Log("Flush and close subkey %s %u", (Status == ERROR_SUCCESS) ? "SUCCESS" : "FAILED", Status);
 439
 440     Status = RegFlushKey(hKey);
 441     Status += RegCloseKey(hKey);
 442     Log("Flush and close key %s %u", (Status == ERROR_SUCCESS) ? "SUCCESS" : "FAILED", Status);
 443
 444     ret = 0;
 445
 446 End:
 447     if (Unload)
 448     {
 449         Status = RegUnLoadKeyA(HKEY_LOCAL_MACHINE, "VTOYNEWSW");
 450         Log("RegUnLoadKey %s %u", (Status == ERROR_SUCCESS) ? "SUCCESS" : "FAILED", Status);
 451     }
 452
 453     if (p)
 454     {
 455         free(p);
 456     }
 457
 458     return ret;
 459 }
 460
 461 static DWORD WINAPI WaitSetupFinishThread(void* Param)
 462 {
 463     CHAR Drv;
 464     BOOL Found = FALSE;
 465     CHAR OfflineRegPath[MAX_PATH];
 466
 467     while (!Found)
 468     {
 469         Sleep(300);
 470         EnumWindows(EnumWindowsProcFunc, (LPARAM)&Found);
 471     }
 472
 473     Log("### Setup finish detected");
 474
 475     // Add HKEY_LOCAL_MACHINE\SOFTWARE\Microsoft\Windows\CurrentVersion\OOBE\BypassNRO
 476     Drv = FindWindowsInstallDstDrive();
 477     Log("Find Windows install drive %d", Drv);
 478     if (Drv)
 479     {
 480         sprintf_s(OfflineRegPath, sizeof(OfflineRegPath), "%C:\\Windows\\system32\\config\\SOFTWARE", Drv);
 481         AddBypassNROReg(OfflineRegPath);
 482     }
 483
 484     return 0;
 485 }
 486
 487 int SetupMonNroStart(const char *isopath)
 488 {
 489     SYSTEMTIME systime;
 490
 491     Log("SetupMonNroStart ...");
 492
 493     FindAllPromptStrings();
 494
 495     if (!g_MuiDataHead)
 496     {
 497         Log("Prompt not found, add default");
 498         g_MuiDataHead = malloc(sizeof(VTOY_MUI_DATA));
 499         if (g_MuiDataHead)
 500         {
 501             wcscpy_s(g_MuiDataHead->Prompt, PROMPT_MAX, L"Windows needs to restart to continue");
 502             g_MuiDataHead->pNext = NULL;
 503         }
 504     }
 505
 506     Log("Wait for setup finish...");
 507     GetSystemTime(&systime);
 508     SystemTimeToFileTime(&systime, &g_SetupMonStartTime);
 509
 510     CreateThread(NULL, 0, WaitSetupFinishThread, NULL, 0, NULL);
 511
 512     return 0;
 513 }