From 83c392a202770914b6fc94c7fbcfbe499bd53b9d Mon Sep 17 00:00:00 2001
From: Nuck-TH <Nuck-Team@yandex.ru>
Date: Thu, 10 Aug 2017 00:10:52 +0700
Subject: [PATCH] Implemented safety(sanity) checks in restoreNAND():    
 *Check for match of file and NAND sizes     *Check for precence of MBR in
 file(this means that file is not encrypted) Added user-friendly message for
 successful finish of restoration. Edited messages a bit in restore function.
 Fixed building on last devkitARM(nand_WriteSectors const buffer in
 declaration). Changed version to 1.6.1 to differ from base.

---
 Makefile             |   2 +-
 arm9/source/main.cpp | 124 ++++++++++++++++++++++++++++---------------
 2 files changed, 82 insertions(+), 44 deletions(-)

diff --git a/Makefile b/Makefile
index 049b3d3..6c0e339 100644
--- a/Makefile
+++ b/Makefile
@@ -10,7 +10,7 @@ include $(DEVKITARM)/ds_rules
 export TARGET	:=	fwTool
 export TOPDIR	:=	$(CURDIR)
 
-export VERSION	:=	1.4.2
+export VERSION	:=	1.6.1
 
 .PHONY: arm7/$(TARGET).elf arm9/$(TARGET).elf
 
diff --git a/arm9/source/main.cpp b/arm9/source/main.cpp
index cf50f00..b00fd23 100644
--- a/arm9/source/main.cpp
+++ b/arm9/source/main.cpp
@@ -14,7 +14,7 @@
 
 extern "C" {
 	bool nand_ReadSectors(sec_t sector, sec_t numSectors,void* buffer);
-	bool nand_WriteSectors(sec_t sector, sec_t numSectors,void* buffer); //!!!
+	bool nand_WriteSectors(sec_t sector, sec_t numSectors,const void* buffer); //!!!
 }
 
 int menuTop = 5, statusTop = 18;
@@ -228,54 +228,92 @@ void restoreNAND() {
 
 	if (!isDSiMode()) {
 		iprintf("Not a DSi or 3ds!\n");
-	} else {
+		return;
+	}
+	
+	FILE *f = fopen(nand_type, "rb");
+	
+	if (f == NULL) {
+		iprintf("Failure opening %s\n", nand_type);
+		return;
+	}
+	
+	//Sanity checks
+	//	Size check
+	fseek(f, 0, SEEK_END);
+	size_t dump_size = ftell(f);
+	if ( dump_size != (sizMB * 1024 * 1024) ) {
+		iprintf("%s and NAND sizes\ndo not match.\nOperation aborted.", nand_type);
+		fclose(f);
+		return;
+	}
+	rewind(f);
+	//	MBR(decrypted image) check
+	//	Taken from TWLtool (https://github.com/WinterMute/twltool)
+	struct {
+		u8 code[446];
+		struct {
+			u8  status;
+			u8  start_chs[3];
+			u8  partition_type;
+			u8  end_chs[3];
+			u32 start_sector;
+			u32 num_sectors;
+		} __attribute__((__packed__)) partition[4];
+		u8 signature[2];
+	} mbr;
+    fread(&mbr, 1, 0x200, f);
+    if(mbr.signature[0] == 0x55 || mbr.signature[1] == 0xAA) {
+		iprintf("Found MBR in %s.\nImage is not encrypted.\nOperation aborted.", nand_type);
+		fclose(f);
+		return;
+	}
+	rewind(f);
+	//Sanity checks end
+	
+	iprintf("Sure? NAND restore is DANGEROUS!");
+	iprintf("START + SELECT confirm\n");
+	iprintf("B to cancel\n");
+	
+	while (1) {
+	    scanKeys();
+		int keys = keysHeld();
+		if ((keys & KEY_START) && (keys & KEY_SELECT))break;
+		if (keys & KEY_B) {
+			clearStatus();
+			fclose(f);
+			return;
+		}
+		swiWaitForVBlank();
+	}
+	
+	clearStatus();
+	
+	iprintf("Reading %s/%s\n\n", dirname, nand_type);
+	size_t i;
+	size_t sectors = 128;
+	size_t blocks = (sizMB * 1024 * 1024) / (sectors * 512);
+	for (i=0; i < blocks; i++) {
 		
-		iprintf("Sure? NAND restore is DANGEROUS!");
-		iprintf("START + SELECT confirm\n");
-		iprintf("B to exit\n");
+		size_t read = fread(firmware_buffer, 1, 512 * sectors, f);
 		
-		while(1){
-		    scanKeys();
-			int keys = keysHeld();
-			if((keys & KEY_START) && (keys & KEY_SELECT))break;
-			if(keys & KEY_B){
-				clearStatus();
-				return;
-			}
-			swiWaitForVBlank();
+		if(read != 512 * sectors) {
+			iprintf("\nError reading SD!\n");
+			break;
 		}
 		
-		clearStatus();
-	
-		FILE *f = fopen(nand_type, "rb");
-
-		if (NULL == f) {
-			iprintf("failure creating %s\n", nand_type);
-		} else {
-			iprintf("Reading %s/%s\n\n", dirname, nand_type);
-			size_t i;
-			size_t sectors = 128;
-			size_t blocks = (sizMB * 1024 * 1024) / (sectors * 512);
-			for (i=0; i < blocks; i++) {
-				
-				size_t read = fread(firmware_buffer, 1, 512 * sectors, f);
-				
-				if(read != 512 * sectors) {
-					iprintf("\nError reading SD!\n");
-					break;
-				}
-				
-				if(!nand_WriteSectors(i * sectors,sectors,firmware_buffer)) {
-					iprintf("\nError writing NAND!\n");
-					break;
-				}
-				
-				iprintf("%d/%d DON'T poweroff!\r", i+1, blocks);
-			}
-			fclose(f);
+		if(!nand_WriteSectors(i * sectors,sectors,firmware_buffer)) {
+			iprintf("\nError writing NAND!\n");
+			break;
 		}
+		
+		iprintf("%d/%d DON'T poweroff!\r", i+1, blocks);
 	}
-
+	fclose(f);
+	
+	clearStatus();
+	iprintf("Restore %s success.\nYou may now Exit and restart\nyour console.",nand_type);
+	
 }
 
 void dumpCID(){