3 files changed, 24 insertions, 40 deletions

diff --git a/twinstall.cpp b/twinstall.cpp
index fa6b8594e..e4b11a7b6 100644..100755
--- a/twinstall.cpp
+++ b/twinstall.cpp
@@ -332,41 +332,12 @@ int TWinstall_zip(const char* path, int* wipe_cache) {
 	if (strlen(path) < 9 || strncmp(path, "/sideload", 9) != 0) {
 		string digest_str;
 		string Full_Filename = path;
-		string digest_file = path;
-		string defmd5file = digest_file + ".md5sum";
-
-		if (TWFunc::Path_Exists(defmd5file)) {
-			digest_file += ".md5sum";
-		}
-		else {
-			digest_file += ".md5";
-		}
 
 		gui_msg("check_for_digest=Checking for Digest file...");
-		if (!TWFunc::Path_Exists(digest_file)) {
-			gui_msg("no_digest=Skipping Digest check: no Digest file found");
-		}
-		else {
-			if (TWFunc::read_file(digest_file, digest_str) != 0) {
-				LOGERR("Skipping MD5 check: MD5 file unreadable\n");
-			}
-			else {
-				twrpDigest *digest = new twrpMD5();
-				if (!twrpDigestDriver::stream_file_to_digest(Full_Filename, digest)) {
-					delete digest;
-					return INSTALL_CORRUPT;
-				}
-				string digest_check = digest->return_digest_string();
-				if (digest_str == digest_check) {
-					gui_msg(Msg("digest_matched=Digest matched for '{1}'.")(path));
-				}
-				else {
+
+		if (!twrpDigestDriver::Check_File_Digest(Full_Filename)) {
 					LOGERR("Aborting zip install: Digest verification failed\n");
-					delete digest;
 					return INSTALL_CORRUPT;
-				}
-				delete digest;
-			}
 		}
 	}
 
diff --git a/twrpDigestDriver.cpp b/twrpDigestDriver.cpp
index c27545c8d..079d0b8f5 100644..100755
--- a/twrpDigestDriver.cpp
+++ b/twrpDigestDriver.cpp
@@ -33,7 +33,7 @@
 #include "twrpDigest/twrpSHA.hpp"
 
 
-bool twrpDigestDriver::Check_Restore_File_Digest(const string& Filename) {
+bool twrpDigestDriver::Check_File_Digest(const string& Filename) {
 	twrpDigest *digest;
 	string digestfile = Filename, file_name = Filename;
 	string digest_str;
@@ -47,18 +47,31 @@ bool twrpDigestDriver::Check_Restore_File_Digest(const string& Filename) {
 		use_sha2 = true;
 	}
 	else {
-		digest = new twrpMD5();
-		digestfile = Filename + ".md5";
-
+		digestfile = Filename + ".sha256";
+		if (TWFunc::Path_Exists(digestfile)) {
+			digest = new twrpSHA256();
+			use_sha2 = true;
+		} else {
+			digest = new twrpMD5();
+			digestfile = Filename + ".md5";
+			if (!TWFunc::Path_Exists(digestfile)) {
+				digestfile = Filename + ".md5sum";
+			}
+		}
 	}
 #else
 	digest = new twrpMD5();
 	digestfile = Filename + ".md5";
+	if (!TWFunc::Path_Exists(digestfile)) {
+		digestfile = Filename + ".md5sum";
+	}
 
 #endif
 
 	if (!TWFunc::Path_Exists(digestfile)) {
-		gui_msg(Msg(msg::kError, "no_digest_found=No digest file found for '{1}'. Please unselect Enable Digest verification to restore.")(Filename));
+		if (Filename.find(".zip") == std::string::npos) {
+			gui_msg(Msg(msg::kError, "no_digest_found=No digest file found for '{1}'. Please unselect Enable Digest verification to restore.")(Filename));
+		}
 		delete digest;
 		return false;
 	}
@@ -80,6 +93,7 @@ bool twrpDigestDriver::Check_Restore_File_Digest(const string& Filename) {
 			LOGINFO("SHA2 Digest: %s  %s\n", digest_str.c_str(), TWFunc::Get_Filename(Filename).c_str());
 		else
 			LOGINFO("MD5 Digest: %s  %s\n", digest_str.c_str(), TWFunc::Get_Filename(Filename).c_str());
+		gui_msg(Msg("digest_matched=Digest matched for '{1}'.")(Filename));
 		delete digest;
 		return true;
 	}
@@ -87,7 +101,6 @@ bool twrpDigestDriver::Check_Restore_File_Digest(const string& Filename) {
 	gui_msg(Msg(msg::kError, "digest_fail_match=Digest failed to match on '{1}'.")(Filename));
 	delete digest;
 	return false;
-
 }
 
 bool twrpDigestDriver::Check_Digest(string Full_Filename) {
@@ -103,13 +116,13 @@ bool twrpDigestDriver::Check_Digest(string Full_Filename) {
 			if (!TWFunc::Path_Exists(split_filename))
 				break;
 				LOGINFO("split_filename: %s\n", split_filename);
-				if (!Check_Restore_File_Digest(split_filename))
+				if (!Check_File_Digest(split_filename))
 					return false;
 				index++;
 		}
 		return true;
 	}
-	return Check_Restore_File_Digest(Full_Filename); // Single file archive
+	return Check_File_Digest(Full_Filename); // Single file archive
 }
 
 bool twrpDigestDriver::Write_Digest(string Full_Filename) {
diff --git a/twrpDigestDriver.hpp b/twrpDigestDriver.hpp
index 6902c8ec6..64e977833 100644..100755
--- a/twrpDigestDriver.hpp
+++ b/twrpDigestDriver.hpp
@@ -24,7 +24,7 @@
 class twrpDigestDriver {
 public:
 
-	static bool Check_Restore_File_Digest(const string& Filename);		//Check the digest of a TWRP partition backup
+	static bool Check_File_Digest(const string& Filename);		//Check the digest of a TWRP partition backup
 	static bool Check_Digest(string Full_Filename);				//Check to make sure the digest is correct
 	static bool Write_Digest(string Full_Filename);				//Write the digest to a file
 	static bool Make_Digest(string Full_Filename);				//Create the digest for a partition backup