1 files changed, 204 insertions, 69 deletions
diff --git a/DcsRe/DcsRe.c b/DcsRe/DcsRe.c
index 7f4096f..aa6f2ed 100644
--- a/DcsRe/DcsRe.c
+++ b/DcsRe/DcsRe.c
@@ -21,11 +21,15 @@ https://opensource.org/licenses/LGPL-3.0
 #include "common/Tcdefs.h"
 
 #ifdef _M_X64
+#define ARCHdot L"x64."
 #define ARCHdotEFI L"x64.efi"
 #else
+#define ARCHdot L"IA32."
 #define ARCHdotEFI L"IA32.efi"
 #endif
 
+CONST CHAR8* g_szMsBootString = "bootmgfw.pdb";
+CONST CHAR16* g_szVcBootString = L"VeraCrypt";
 
 //////////////////////////////////////////////////////////////////////////
 // Menu
@@ -48,48 +52,56 @@ SelectEfiVolume()
 	EFI_FILE     *file;
 	EFI_FILE     **efiVolumes;
 	UINTN        efiVolumesCount = 0;
-	EFI_HANDLE   startHandle;
 	if (EfiBootVolume != NULL) return;
-	res = EfiGetStartDevice(&startHandle);
-	if (EFI_ERROR(res)) {
-		ERR_PRINT(L"GetStartDevice %r", res);
-		return;
-	}
+
 	efiVolumes = MEM_ALLOC(sizeof(EFI_FILE*) * gFSCount);
 	for (i = 0; i < gFSCount; ++i) {
+		if (gFSHandles[i] == gFileRootHandle)
+			continue;
 		res = FileOpenRoot(gFSHandles[i], &file);
-		if(EFI_ERROR(res)) continue;
-		if (!EFI_ERROR(FileExist(file, L"EFI\\Boot\\boot" ARCHdotEFI))) {
+		if(EFI_ERROR(res)) { ERR_PRINT(L"FileOpenRoot %r\n", res); continue;}
+		if (	!EFI_ERROR(FileExist(file, L"EFI\\Boot\\boot" ARCHdotEFI))
+			||	!EFI_ERROR(FileExist(file, L"EFI\\Microsoft\\Boot\\bootmgfw.efi"))
+			||	!EFI_ERROR(FileExist(file, L"EFI\\Microsoft\\Boot\\bootmgfw_ms.vc"))
+			) 
+		{
 			efiVolumesCount++;
 			efiVolumes[i] = file;
-			if (gFSHandles[i] != startHandle) {
-				EfiBootVolumeIndex = i;
-				EfiBootVolume = file;
-			}
+			EfiBootVolumeIndex = i;
+			EfiBootVolume = file;
 		}	else {
 			FileClose(file);
 		}
 	}
-
-	for (i = 0; i < gFSCount; ++i) {
-		OUT_PRINT(L"%H%d)%N ", i);
-		if (efiVolumes[i] != NULL) {
-			if (gFSHandles[i] == startHandle) {
-				OUT_PRINT(L"%V [Boot Rescue] %N");
-			}
-			else {
+	
+	if (efiVolumesCount > 1)
+	{
+		for (i = 0; i < gFSCount; ++i) {
+			OUT_PRINT(L"%H%d)%N ", i);
+			if (efiVolumes[i] != NULL) {
 				OUT_PRINT(L"%V [Boot] %N");
 			}
+			EfiPrintDevicePath(gFSHandles[i]);
+			OUT_PRINT(L"\n");
 		}
-		EfiPrintDevicePath(gFSHandles[i]);
-		OUT_PRINT(L"\n");
-	}
 
-	do {
-		EfiBootVolumeIndex = AskUINTN("Select EFI boot volume:", EfiBootVolumeIndex);
-		if (EfiBootVolumeIndex >= gFSCount) continue;
-		EfiBootVolume = efiVolumes[EfiBootVolumeIndex];
-	} while (EfiBootVolume == NULL);
+		do {
+			EfiBootVolumeIndex = AskUINTN("Select EFI boot volume:", EfiBootVolumeIndex);
+			if (EfiBootVolumeIndex >= gFSCount) continue;
+			EfiBootVolume = efiVolumes[EfiBootVolumeIndex];
+		} while (EfiBootVolume == NULL);
+		
+		/* free unused descriptors */
+		for (i = 0; i < gFSCount; ++i) {
+			if (efiVolumes[i] != NULL && efiVolumes[i] != EfiBootVolume) {
+				FileClose(efiVolumes[i]);
+			}
+		}
+
+		OUT_PRINT (L"\n");
+	}
+	
+	
 	MEM_FREE(efiVolumes);
 }
 
@@ -113,6 +125,54 @@ ActionDcsBoot(IN VOID* ctx) {
 	return EfiExec(gFSHandles[EfiBootVolumeIndex], L"EFI\\VeraCrypt\\DcsBoot.efi");
 }
 
+EFI_STATUS
+ActionWindowsBoot(IN VOID* ctx) {
+	if (AskConfirm("If Windows is encrypted, Windows original loader will fail to start.\r\nDo you want to continue? [N]", 1))
+	{
+		SelectEfiVolume();
+		if (EfiBootVolume == NULL) return EFI_NOT_READY;
+		if (!EFI_ERROR(FileExist(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw_ms.vc")))
+			return EfiExec(gFSHandles[EfiBootVolumeIndex], L"EFI\\Microsoft\\Boot\\bootmgfw_ms.vc");
+		else
+		{
+			if (!EFI_ERROR(FileExist(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi")))
+			{
+				/* check if it is Microsoft one */
+				UINT8*      fileData = NULL;
+				UINTN       fileSize = 0;
+				BOOLEAN		bFound = FALSE;
+				if (!EFI_ERROR(FileLoad(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi", &fileData, &fileSize)))
+				{
+					if ((fileSize > 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szMsBootString, AsciiStrLen(g_szMsBootString))))
+					{
+						bFound = TRUE;
+					}
+				}
+				
+				MEM_FREE(fileData);
+				
+				if (bFound)
+					return EfiExec(gFSHandles[EfiBootVolumeIndex], L"EFI\\Microsoft\\Boot\\bootmgfw.efi");
+			}
+			
+			/* copy our backup copy and then boot from it*/
+			if (!EFI_ERROR(FileExist(NULL, L"\\EFI\\Boot\\original_boot" ARCHdot L"vc_backup")))
+			{
+				if (!EFI_ERROR(FileCopy(NULL, L"\\EFI\\Boot\\original_boot" ARCHdot L"vc_backup", EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw_ms.vc", 1024 * 1024)))
+				{
+					return EfiExec(gFSHandles[EfiBootVolumeIndex], L"EFI\\Microsoft\\Boot\\bootmgfw_ms.vc");
+				}
+			}
+
+			ERR_PRINT(L"Could not find the original Windows loader\r\n");
+			
+			return EFI_NOT_READY;
+		}
+	}
+	else
+		return EFI_SUCCESS;
+}
+
 CHAR16* DcsBootBins[] = {
 	L"EFI\\VeraCrypt\\DcsBoot.efi",
 	L"EFI\\VeraCrypt\\DcsInt.dcs",
@@ -130,11 +190,73 @@ ActionRestoreDcsLoader(IN VOID* ctx) {
 	UINTN i;
 	SelectEfiVolume();
 	if (EfiBootVolume == NULL) return EFI_NOT_READY;
+	
+	DirectoryCreate (EfiBootVolume, L"EFI\\VeraCrypt");
+	
 	for (i = 0; i < sizeof(DcsBootBins) / sizeof(CHAR16*); ++i) {
 		res = FileCopy(NULL, DcsBootBins[i], EfiBootVolume, DcsBootBins[i], 1024 * 1024);
 		if (EFI_ERROR(res)) return res;
 	}
-	return res;
+	/* restore standard boot file */
+	if (!EFI_ERROR(FileExist(EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI)))
+	{
+		/* check if it is Microsoft one or ours */
+		UINT8*      fileData = NULL;
+		UINTN       fileSize = 0;
+		res = EFI_SUCCESS;
+		if (!EFI_ERROR(FileLoad(EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, &fileData, &fileSize)))
+		{
+			if ((fileSize > 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szMsBootString, AsciiStrLen(g_szMsBootString))))
+			{
+				res = FileCopy(EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, EfiBootVolume, L"\\EFI\\Boot\\original_boot" ARCHdot L"vc_backup", 1024 * 1024);
+				if (!EFI_ERROR(res))
+					res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, 1024 * 1024);				
+			}
+			else if ((fileSize <= 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szVcBootString, StrLen (g_szVcBootString) * 2)))
+			{
+				res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, 1024 * 1024);
+			}
+			MEM_FREE(fileData);
+			
+			if (EFI_ERROR(res)) return res;
+		}		
+	}
+	else if (!EFI_ERROR(FileExist(EfiBootVolume, L"\\EFI\\Boot\\original_boot" ARCHdot L"vc_backup")))
+	{
+		res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"EFI\\Boot\\boot" ARCHdotEFI, 1024 * 1024);
+		if (EFI_ERROR(res)) return res;
+	}
+	
+	if (!EFI_ERROR(FileExist(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi")))
+	{
+		/* check if it is Microsoft one */
+		UINT8*      fileData = NULL;
+		UINTN       fileSize = 0;
+		res = EFI_SUCCESS;
+		if (!EFI_ERROR(FileLoad(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi", &fileData, &fileSize)))
+		{
+			if ((fileSize > 32768) && !EFI_ERROR(MemoryHasPattern(fileData, fileSize, g_szMsBootString, AsciiStrLen(g_szMsBootString))))
+			{
+				res = FileCopy(EfiBootVolume, L"EFI\\Microsoft\\Boot\\bootmgfw.efi", EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw_ms.vc", 1024 * 1024);				
+			}
+
+			MEM_FREE(fileData);
+			
+			if (EFI_ERROR(res)) return res;
+		}
+
+		res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw.efi", 1024 * 1024);
+		if (EFI_ERROR(res)) return res;
+	}
+	else if (!EFI_ERROR(FileExist(EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw_ms.vc")))		
+	{
+		res = FileCopy(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi", EfiBootVolume, L"\\EFI\\Microsoft\\Boot\\bootmgfw.efi", 1024 * 1024);
+		if (EFI_ERROR(res)) return res;
+	}
+	
+	OUT_PRINT (L"\nVeraCrypt Loader restored to disk successfully\n\n");
+	
+	return EFI_SUCCESS;
 }
 
 CHAR16* sDcsBootEfi = L"EFI\\VeraCrypt\\DcsBoot.efi";
@@ -236,54 +358,67 @@ DcsReMain(
       ERR_PRINT(L"InitFS %r\n", res);
 		return res;
    }
+   
+	if (!EFI_ERROR(DirectoryExists(NULL, L"EFI\\VeraCrypt")))
+	{
+		item = DcsMenuAppend(NULL, L"Decrypt OS", 'd', ActionDecryptOS, NULL);
+		gMenu = item;
+		item = DcsMenuAppend(item, L"Restore VeraCrypt loader to boot menu", 'm', ActionRestoreDcsBootMenu, NULL);
+		item = DcsMenuAppend(item, L"Remove VeraCrypt loader from boot menu", 'z' , ActionRemoveDcsBootMenu, NULL);
+
+		if (!EFI_ERROR(FileExist(NULL, L"EFI\\VeraCrypt\\DcsProp"))) {
+			item = DcsMenuAppend(item, L"Restore VeraCrypt loader configuration to system disk", 'c', ActionRestoreDcsProp, NULL);
+		}
 
-	item = DcsMenuAppend(NULL, L"Decrypt OS", 'd', ActionDecryptOS, NULL);
-	gMenu = item;
-	item = DcsMenuAppend(item, L"Restore VeraCrypt loader to boot menu", 'm', ActionRestoreDcsBootMenu, NULL);
-	item = DcsMenuAppend(item, L"Remove VeraCrypt loader from boot menu", 'z' , ActionRemoveDcsBootMenu, NULL);
-
-	if (!EFI_ERROR(FileExist(NULL, L"EFI\\VeraCrypt\\DcsProp"))) {
-		item = DcsMenuAppend(item, L"Restore VeraCrypt loader configuration to system disk", 'c', ActionRestoreDcsProp, NULL);
-	}
-
-	if (!EFI_ERROR(FileExist(NULL, L"EFI\\VeraCrypt\\svh_bak"))) {
-		item = DcsMenuAppend(item, L"Restore OS header keys", 'k', ActionRestoreHeader, NULL);
-	}
+		if (!EFI_ERROR(FileExist(NULL, L"EFI\\VeraCrypt\\svh_bak"))) {
+			item = DcsMenuAppend(item, L"Restore OS header keys", 'k', ActionRestoreHeader, NULL);
+		}
 
-	if (!EFI_ERROR(FileExist(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi"))) {
-		item = DcsMenuAppend(item, L"Restore VeraCrypt loader binaries to system disk", 'r', ActionRestoreDcsLoader, NULL);
-		item = DcsMenuAppend(item, L"Boot VeraCrypt loader from rescue disk", 'v', ActionDcsBoot, NULL);
-	}
+		if (!EFI_ERROR(FileExist(NULL, L"EFI\\VeraCrypt\\DcsBoot.efi"))) {
+			item = DcsMenuAppend(item, L"Restore VeraCrypt loader binaries to system disk", 'r', ActionRestoreDcsLoader, NULL);
+			item = DcsMenuAppend(item, L"Boot VeraCrypt loader from rescue disk", 'v', ActionDcsBoot, NULL);
+		}
+		
+		item = DcsMenuAppend(item, L"Boot Original Windows Loader", 'o', ActionWindowsBoot, NULL);
 
-	if (!EFI_ERROR(FileExist(NULL, L"EFI\\Boot\\WinPE_boot" ARCHdotEFI))) {
-		item = DcsMenuAppend(item, L"Boot Windows PE from rescue disk", 'w', ActionBootWinPE, NULL);
-	}
+		if (!EFI_ERROR(FileExist(NULL, L"EFI\\Boot\\WinPE_boot" ARCHdotEFI))) {
+			item = DcsMenuAppend(item, L"Boot Windows PE from rescue disk", 'w', ActionBootWinPE, NULL);
+		}
 
-	if (!EFI_ERROR(FileExist(NULL, L"EFI\\Shell\\Shell.efi"))) {
-		item = DcsMenuAppend(item, L"Boot Shell.efi from rescue disk", 's', ActionShell, NULL);
-	}
+		if (!EFI_ERROR(FileExist(NULL, L"EFI\\Shell\\Shell.efi"))) {
+			item = DcsMenuAppend(item, L"Boot Shell.efi from rescue disk", 's', ActionShell, NULL);
+		}
 
-	item = DcsMenuAppend(item, L"Help", 'h', ActionHelp, NULL);
-	item = DcsMenuAppend(item, L"Exit", 'e', ActionExit, NULL);
-	OUT_PRINT(L"%V%a rescue disk %a%N\n", TC_APP_NAME, VERSION_STRING);
-	gBS->SetWatchdogTimer(0, 0, 0, NULL);
-	do {
-		DcsMenuPrint(gMenu);
-		item = NULL;
-		key.UnicodeChar = 0;
-		while (item == NULL) {
-			item = gMenu;
-			key = GetKey();
-			while (item != NULL) {
-				if (item->Select == key.UnicodeChar) break;
-				item = item->Next;
+		item = DcsMenuAppend(item, L"Help", 'h', ActionHelp, NULL);
+		item = DcsMenuAppend(item, L"Exit", 'e', ActionExit, NULL);
+		OUT_PRINT(L"%V%a rescue disk %a%N\n", TC_APP_NAME, VERSION_STRING);
+		gBS->SetWatchdogTimer(0, 0, 0, NULL);
+		do {
+			DcsMenuPrint(gMenu);
+			item = NULL;
+			key.UnicodeChar = 0;
+			while (item == NULL) {
+				item = gMenu;
+				key = GetKey();
+				while (item != NULL) {
+					if (item->Select == key.UnicodeChar) break;
+					item = item->Next;
+				}
 			}
-		}
-		OUT_PRINT(L"%c\n",key.UnicodeChar);
-		res = item->Action(item->Context);
+			OUT_PRINT(L"%c\n",key.UnicodeChar);
+			res = item->Action(item->Context);
+			if (EFI_ERROR(res)) {
+				ERR_PRINT(L"%r\n", res);
+			}
+		} while (gContiniue);
+	}
+	else
+	{
+		/* No VeraCrypt folder. Boot directly from the hard drive */
+		res = ActionDcsBoot (NULL);
 		if (EFI_ERROR(res)) {
 			ERR_PRINT(L"%r\n", res);
 		}
-	} while (gContiniue);
+	}
 	return EFI_INVALID_PARAMETER;
 }