1 files changed, 178 insertions, 104 deletions

diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index 320944c3f..96e30025e 100644
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -27,6 +27,7 @@
 #include <time.h>
 #include <errno.h>
 #include <fcntl.h>
+#include <zlib.h>
 #include <iostream>
 #include <iomanip>
 #include <sys/wait.h>
@@ -44,6 +45,7 @@
 #include "tw_atomic.hpp"
 #include "gui/gui.hpp"
 #include "progresstracking.hpp"
+#include "adbbu/libtwadbbu.hpp"
 
 #ifdef TW_HAS_MTP
 #include "mtp/mtp_MtpServer.hpp"
@@ -68,7 +70,6 @@ TWPartitionManager::TWPartitionManager(void) {
 	mtp_was_enabled = false;
 	mtp_write_fd = -1;
 	stop_backup.set_value(0);
-	tar_fork_pid = 0;
 }
 
 int TWPartitionManager::Process_Fstab(string Fstab_Filename, bool Display_Error) {
@@ -472,6 +473,7 @@ int TWPartitionManager::Check_Backup_Name(bool Display_Error) {
 	if (TWFunc::Path_Exists(tw_image_dir)) {
 		if (Display_Error)
 			gui_err("backup_name_exists=A backup with that name already exists!");
+
 		return -4;
 	}
 	// No problems found, return 0
@@ -481,7 +483,7 @@ int TWPartitionManager::Check_Backup_Name(bool Display_Error) {
 bool TWPartitionManager::Make_MD5(bool generate_md5, string Backup_Folder, string Backup_Filename)
 {
 	string command;
-	string Full_File = Backup_Folder + Backup_Filename;
+	string Full_File = Backup_Folder + "/" + Backup_Filename;
 	string result;
 	twrpDigest md5sum;
 
@@ -490,9 +492,8 @@ bool TWPartitionManager::Make_MD5(bool generate_md5, string Backup_Folder, strin
 
 	TWFunc::GUI_Operation_Text(TW_GENERATE_MD5_TEXT, gui_parse_text("{@generating_md51}"));
 	gui_msg("generating_md52= * Generating md5...");
-
 	if (TWFunc::Path_Exists(Full_File)) {
-		md5sum.setfn(Backup_Folder + Backup_Filename);
+		md5sum.setfn(Full_File);
 		if (md5sum.computeMD5() == 0)
 			if (md5sum.write_md5digest() == 0)
 				gui_msg("md5_created= * MD5 Created.");
@@ -533,13 +534,13 @@ bool TWPartitionManager::Make_MD5(bool generate_md5, string Backup_Folder, strin
 	return true;
 }
 
-bool TWPartitionManager::Backup_Partition(TWPartition* Part, const string& Backup_Folder, bool generate_md5, unsigned long *img_time, unsigned long *file_time, ProgressTracking *progress) {
-	time_t start, stop;
-	int use_compression;
 
-	string backup_log = Backup_Folder + "recovery.log";
+bool TWPartitionManager::Backup_Partition(PartitionSettings *part_settings) {
+	time_t start, stop;
+	int use_compression, adb_control_bu_fd;
+	string backup_log = part_settings->Backup_Folder + "/recovery.log";
 
-	if (Part == NULL)
+	if (part_settings->Part == NULL)
 		return true;
 
 	DataManager::GetValue(TW_USE_COMPRESSION_VAR, use_compression);
@@ -547,25 +548,28 @@ bool TWPartitionManager::Backup_Partition(TWPartition* Part, const string& Backu
 	TWFunc::SetPerformanceMode(true);
 	time(&start);
 
-	if (Part->Backup(Backup_Folder, tar_fork_pid, progress)) {
+	part_settings->Backup_FileName = part_settings->Part->Backup_Name + "." +  part_settings->Part->Current_File_System + ".win";
+	if (part_settings->Part->Backup(part_settings, &tar_fork_pid)) {
 		bool md5Success = false;
-		if (Part->Has_SubPartition) {
+		if (part_settings->Part->Has_SubPartition) {
 			std::vector<TWPartition*>::iterator subpart;
 
 			for (subpart = Partitions.begin(); subpart != Partitions.end(); subpart++) {
-				if ((*subpart)->Can_Be_Backed_Up && (*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == Part->Mount_Point) {
-					if (!(*subpart)->Backup(Backup_Folder, tar_fork_pid, progress)) {
+				if ((*subpart)->Can_Be_Backed_Up && (*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == part_settings->Part->Mount_Point) {
+					if (!(*subpart)->Backup(part_settings, &tar_fork_pid)) {
 						TWFunc::SetPerformanceMode(false);
-						Clean_Backup_Folder(Backup_Folder);
+						Clean_Backup_Folder(part_settings->Backup_Folder);
 						TWFunc::copy_file("/tmp/recovery.log", backup_log, 0644);
 						tw_set_default_metadata(backup_log.c_str());
 						return false;
 					}
 					sync();
 					sync();
-					if (!Make_MD5(generate_md5, Backup_Folder, (*subpart)->Backup_FileName)) {
-						TWFunc::SetPerformanceMode(false);
-						return false;
+					if (!part_settings->adbbackup) {
+						if (!Make_MD5(part_settings->generate_md5, part_settings->Backup_Folder, (*subpart)->Backup_FileName)) {
+							TWFunc::SetPerformanceMode(false);
+							return false;
+						}
 					}
 				}
 			}
@@ -573,17 +577,23 @@ bool TWPartitionManager::Backup_Partition(TWPartition* Part, const string& Backu
 		time(&stop);
 		int backup_time = (int) difftime(stop, start);
 		LOGINFO("Partition Backup time: %d\n", backup_time);
-		if (Part->Backup_Method == 1) {
-			*file_time += backup_time;
+		if (part_settings->Part->Backup_Method == BM_FILES) {
+			part_settings->file_time += backup_time;
 		} else {
-			*img_time += backup_time;
+			part_settings->img_time += backup_time;
+
 		}
 
-		md5Success = Make_MD5(generate_md5, Backup_Folder, Part->Backup_FileName);
+		if (!part_settings->adbbackup) {
+			md5Success = Make_MD5(part_settings->generate_md5, part_settings->Backup_Folder, part_settings->Part->Backup_FileName);
+		}
+		else 
+			md5Success = true;
 		TWFunc::SetPerformanceMode(false);
+
 		return md5Success;
 	} else {
-		Clean_Backup_Folder(Backup_Folder);
+		Clean_Backup_Folder(part_settings->Backup_Folder);
 		TWFunc::copy_file("/tmp/recovery.log", backup_log, 0644);
 		tw_set_default_metadata(backup_log.c_str());
 		TWFunc::SetPerformanceMode(false);
@@ -648,12 +658,12 @@ int TWPartitionManager::Cancel_Backup() {
 	return 0;
 }
 
-int TWPartitionManager::Run_Backup(void) {
-	int check, do_md5, partition_count = 0, disable_free_space_check = 0;
-	string Backup_Folder, Backup_Name, Full_Backup_Path, Backup_List, backup_path;
-	unsigned long long total_bytes = 0, file_bytes = 0, img_bytes = 0, free_space = 0, img_bytes_remaining, file_bytes_remaining, subpart_size;
-	unsigned long img_time = 0, file_time = 0;
-	TWPartition* backup_part = NULL;
+int TWPartitionManager::Run_Backup(bool adbbackup) {
+	PartitionSettings part_settings;
+	int check, partition_count = 0, disable_free_space_check = 0, do_md5 = 0;
+	int gui_adb_backup;
+	string Backup_Name, Backup_List, backup_path;
+	unsigned long long total_bytes = 0, free_space = 0, subpart_size;
 	TWPartition* storage = NULL;
 	std::vector<TWPartition*>::iterator subpart;
 	struct tm *t;
@@ -663,6 +673,19 @@ int TWPartitionManager::Run_Backup(void) {
 	seconds = time(0);
 	t = localtime(&seconds);
 
+	part_settings.img_bytes_remaining = 0;
+	part_settings.file_bytes_remaining = 0;
+	part_settings.img_time = 0;
+	part_settings.file_time = 0;
+	part_settings.img_bytes = 0;
+	part_settings.file_bytes = 0;
+	part_settings.PM_Method = PM_BACKUP;
+
+	DataManager::GetValue("tw_enable_adb_backup", gui_adb_backup);
+	if (gui_adb_backup == true)
+		adbbackup = true;
+
+	part_settings.adbbackup = adbbackup;
 	time(&total_start);
 
 	Update_System_Details();
@@ -670,23 +693,26 @@ int TWPartitionManager::Run_Backup(void) {
 	if (!Mount_Current_Storage(true))
 		return false;
 
+
 	DataManager::GetValue(TW_SKIP_MD5_GENERATE_VAR, do_md5);
 	if (do_md5 == 0)
-		do_md5 = true;
+		part_settings.generate_md5 = true;
 	else
-		do_md5 = false;
+		part_settings.generate_md5 = false;
 
-	DataManager::GetValue(TW_BACKUPS_FOLDER_VAR, Backup_Folder);
-	DataManager::GetValue(TW_BACKUP_NAME, Backup_Name);
-	if (Backup_Name == gui_lookup("curr_date", "(Current Date)")) {
-		Backup_Name = TWFunc::Get_Current_Date();
-	} else if (Backup_Name == gui_lookup("auto_generate", "(Auto Generate)") || Backup_Name == "0" || Backup_Name.empty()) {
+	DataManager::GetValue(TW_BACKUPS_FOLDER_VAR, part_settings.Backup_Folder);
+	DataManager::GetValue(TW_BACKUP_NAME, part_settings.Backup_Name);
+	if (part_settings.Backup_Name == gui_lookup("curr_date", "(Current Date)")) {
+		part_settings.Backup_Name = TWFunc::Get_Current_Date();
+	} else if (part_settings.Backup_Name == gui_lookup("auto_generate", "(Auto Generate)") || part_settings.Backup_Name == "0" || part_settings.Backup_Name.empty()) {
 		TWFunc::Auto_Generate_Backup_Name();
-		DataManager::GetValue(TW_BACKUP_NAME, Backup_Name);
+		DataManager::GetValue(TW_BACKUP_NAME, part_settings.Backup_Name);
 	}
-	LOGINFO("Backup Name is: '%s'\n", Backup_Name.c_str());
-	Full_Backup_Path = Backup_Folder + "/" + Backup_Name + "/";
-	LOGINFO("Full_Backup_Path is: '%s'\n", Full_Backup_Path.c_str());
+
+	LOGINFO("Backup Name is: '%s'\n", part_settings.Backup_Name.c_str());
+	part_settings.Full_Backup_Path = part_settings.Backup_Folder + "/" + part_settings.Backup_Name + "/";
+
+	LOGINFO("Full_Backup_Path is: '%s'\n", part_settings.Full_Backup_Path.c_str());
 
 	LOGINFO("Calculating backup details...\n");
 	DataManager::GetValue("tw_backup_list", Backup_List);
@@ -694,23 +720,23 @@ int TWPartitionManager::Run_Backup(void) {
 		end_pos = Backup_List.find(";", start_pos);
 		while (end_pos != string::npos && start_pos < Backup_List.size()) {
 			backup_path = Backup_List.substr(start_pos, end_pos - start_pos);
-			backup_part = Find_Partition_By_Path(backup_path);
-			if (backup_part != NULL) {
+			part_settings.Part = Find_Partition_By_Path(backup_path);
+			if (part_settings.Part != NULL) {
 				partition_count++;
-				if (backup_part->Backup_Method == 1)
-					file_bytes += backup_part->Backup_Size;
+				if (part_settings.Part->Backup_Method == BM_FILES)
+					part_settings.file_bytes += part_settings.Part->Backup_Size;
 				else
-					img_bytes += backup_part->Backup_Size;
-				if (backup_part->Has_SubPartition) {
+					part_settings.img_bytes += part_settings.Part->Backup_Size;
+				if (part_settings.Part->Has_SubPartition) {
 					std::vector<TWPartition*>::iterator subpart;
 
 					for (subpart = Partitions.begin(); subpart != Partitions.end(); subpart++) {
-						if ((*subpart)->Can_Be_Backed_Up && (*subpart)->Is_Present && (*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == backup_part->Mount_Point) {
+						if ((*subpart)->Can_Be_Backed_Up && (*subpart)->Is_Present && (*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == part_settings.Part->Mount_Point) {
 							partition_count++;
-							if ((*subpart)->Backup_Method == 1)
-								file_bytes += (*subpart)->Backup_Size;
+							if ((*subpart)->Backup_Method == BM_FILES)
+								part_settings.file_bytes += (*subpart)->Backup_Size;
 							else
-								img_bytes += (*subpart)->Backup_Size;
+								part_settings.img_bytes += (*subpart)->Backup_Size;
 						}
 					}
 				}
@@ -726,8 +752,15 @@ int TWPartitionManager::Run_Backup(void) {
 		gui_msg("no_partition_selected=No partitions selected for backup.");
 		return false;
 	}
-	total_bytes = file_bytes + img_bytes;
+	if (adbbackup) {
+		if (twadbbu::Write_ADB_Stream_Header(partition_count) == false) {
+			return false;
+		}
+	}
+	total_bytes = part_settings.file_bytes + part_settings.img_bytes;
 	ProgressTracking progress(total_bytes);
+	part_settings.progress = &progress;
+
 	gui_msg(Msg("total_partitions_backup= * Total number of partitions to back up: {1}")(partition_count));
 	gui_msg(Msg("total_backup_size= * Total size of all data: {1}MB")(total_bytes / 1024 / 1024));
 	storage = Find_Partition_By_Path(DataManager::GetCurrentStoragePath());
@@ -740,6 +773,10 @@ int TWPartitionManager::Run_Backup(void) {
 	}
 
 	DataManager::GetValue("tw_disable_free_space", disable_free_space_check);
+
+	if (adbbackup)
+		disable_free_space_check = true;
+
 	if (!disable_free_space_check) {
 		if (free_space - (32 * 1024 * 1024) < total_bytes) {
 			// We require an extra 32MB just in case
@@ -747,12 +784,12 @@ int TWPartitionManager::Run_Backup(void) {
 			return false;
 		}
 	}
-	img_bytes_remaining = img_bytes;
-	file_bytes_remaining = file_bytes;
+	part_settings.img_bytes_remaining = part_settings.img_bytes;
+	part_settings.file_bytes_remaining = part_settings.file_bytes;
 
 	gui_msg("backup_started=[BACKUP STARTED]");
-	gui_msg(Msg("backup_folder= * Backup Folder: {1}")(Full_Backup_Path));
-	if (!TWFunc::Recursive_Mkdir(Full_Backup_Path)) {
+	gui_msg(Msg("backup_folder= * Backup Folder: {1}")(part_settings.Full_Backup_Path));
+	if (!TWFunc::Recursive_Mkdir(part_settings.Full_Backup_Path)) {
 		gui_err("fail_backup_folder=Failed to make backup folder.");
 		return false;
 	}
@@ -765,9 +802,10 @@ int TWPartitionManager::Run_Backup(void) {
 		if (stop_backup.get_value() != 0)
 			return -1;
 		backup_path = Backup_List.substr(start_pos, end_pos - start_pos);
-		backup_part = Find_Partition_By_Path(backup_path);
-		if (backup_part != NULL) {
-			if (!Backup_Partition(backup_part, Full_Backup_Path, do_md5, &img_time, &file_time, &progress))
+		part_settings.Part = Find_Partition_By_Path(backup_path);
+		if (part_settings.Part != NULL) {
+			if (!Backup_Partition(&part_settings))
+
 				return false;
 		} else {
 			gui_msg(Msg(msg::kError, "unable_to_locate_partition=Unable to locate '{1}' partition for backup calculations.")(backup_path));
@@ -777,26 +815,30 @@ int TWPartitionManager::Run_Backup(void) {
 	}
 
 	// Average BPS
-	if (img_time == 0)
-		img_time = 1;
-	if (file_time == 0)
-		file_time = 1;
-	int img_bps = (int)img_bytes / (int)img_time;
-	unsigned long long file_bps = file_bytes / (int)file_time;
-
-	if (file_bytes != 0)
+	if (part_settings.img_time == 0)
+		part_settings.img_time = 1;
+	if (part_settings.file_time == 0)
+		part_settings.file_time = 1;
+	int img_bps = (int)part_settings.img_bytes / (int)part_settings.img_time;
+	unsigned long long file_bps = part_settings.file_bytes / (int)part_settings.file_time;
+
+	if (part_settings.file_bytes != 0)
 		gui_msg(Msg("avg_backup_fs=Average backup rate for file systems: {1} MB/sec")(file_bps / (1024 * 1024)));
-	if (img_bytes != 0)
+	if (part_settings.img_bytes != 0)
 		gui_msg(Msg("avg_backup_img=Average backup rate for imaged drives: {1} MB/sec")(img_bps / (1024 * 1024)));
 
 	time(&total_stop);
 	int total_time = (int) difftime(total_stop, total_start);
-	uint64_t actual_backup_size = du.Get_Folder_Size(Full_Backup_Path);
+
+	uint64_t actual_backup_size;
+	if (!adbbackup)
+		actual_backup_size = du.Get_Folder_Size(part_settings.Full_Backup_Path);
+	else
+		actual_backup_size = part_settings.file_bytes + part_settings.img_bytes;
 	actual_backup_size /= (1024LLU * 1024LLU);
 
-	int prev_img_bps, use_compression;
-	unsigned long long prev_file_bps;
-	DataManager::GetValue(TW_BACKUP_AVG_IMG_RATE, prev_img_bps);
+	int prev_img_bps = 0, use_compression = 0;
+	unsigned long long prev_file_bps = 0;
 	img_bps += (prev_img_bps * 4);
 	img_bps /= 5;
 
@@ -818,27 +860,40 @@ int TWPartitionManager::Run_Backup(void) {
 	Update_System_Details();
 	UnMount_Main_Partitions();
 	gui_msg(Msg(msg::kHighlight, "backup_completed=[BACKUP COMPLETED IN {1} SECONDS]")(total_time)); // the end
-	string backup_log = Full_Backup_Path + "recovery.log";
+	string backup_log = part_settings.Full_Backup_Path + "recovery.log";
 	TWFunc::copy_file("/tmp/recovery.log", backup_log, 0644);
 	tw_set_default_metadata(backup_log.c_str());
+
+	if (part_settings.adbbackup) {
+		if (twadbbu::Write_ADB_Stream_Trailer() == false) {
+			return false;
+		}
+	}
+	part_settings.adbbackup = false;
+	DataManager::SetValue("tw_enable_adb_backup", 0);
+
 	return true;
 }
 
-bool TWPartitionManager::Restore_Partition(TWPartition* Part, const string& Restore_Name, ProgressTracking *progress) {
+bool TWPartitionManager::Restore_Partition(PartitionSettings *part_settings) {
 	time_t Start, Stop;
+
 	TWFunc::SetPerformanceMode(true);
+
 	time(&Start);
 
-	if (!Part->Restore(Restore_Name, progress)) {
+
+	if (!part_settings->Part->Restore(part_settings)) {
 		TWFunc::SetPerformanceMode(false);
 		return false;
 	}
-	if (Part->Has_SubPartition) {
+	if (part_settings->Part->Has_SubPartition) {
 		std::vector<TWPartition*>::iterator subpart;
 
 		for (subpart = Partitions.begin(); subpart != Partitions.end(); subpart++) {
-			if ((*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == Part->Mount_Point) {
-				if (!(*subpart)->Restore(Restore_Name, progress)) {
+
+			if ((*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == part_settings->Part->Mount_Point) {
+				if (!(*subpart)->Restore(part_settings)) {
 					TWFunc::SetPerformanceMode(false);
 					return false;
 				}
@@ -847,18 +902,27 @@ bool TWPartitionManager::Restore_Partition(TWPartition* Part, const string& Rest
 	}
 	time(&Stop);
 	TWFunc::SetPerformanceMode(false);
-	gui_msg(Msg("restore_part_done=[{1} done ({2} seconds)]")(Part->Backup_Display_Name)((int)difftime(Stop, Start)));
+	gui_msg(Msg("restore_part_done=[{1} done ({2} seconds)]")(part_settings->Part->Backup_Display_Name)((int)difftime(Stop, Start)));
+
 	return true;
 }
 
 int TWPartitionManager::Run_Restore(const string& Restore_Name) {
+	PartitionSettings part_settings;
 	int check_md5, check, partition_count = 0;
 	TWPartition* restore_part = NULL;
+
 	time_t rStart, rStop;
 	time(&rStart);
 	string Restore_List, restore_path;
 	size_t start_pos = 0, end_pos;
-	unsigned long long total_restore_size = 0, already_restored_size = 0;
+
+	part_settings.Restore_Name = Restore_Name;
+	part_settings.Part = NULL;
+	part_settings.partition_count = 0;
+	part_settings.total_restore_size = 0;
+	part_settings.adbbackup = false;
+	part_settings.PM_Method = PM_RESTORE;
 
 	gui_msg("restore_started=[RESTORE STARTED]");
 	gui_msg(Msg("restore_folder=Restore folder: '{1}'")(Restore_Name));
@@ -876,28 +940,31 @@ int TWPartitionManager::Run_Restore(const string& Restore_Name) {
 	}
 	gui_msg("calc_restore=Calculating restore details...");
 	DataManager::GetValue("tw_restore_selected", Restore_List);
+
 	if (!Restore_List.empty()) {
 		end_pos = Restore_List.find(";", start_pos);
 		while (end_pos != string::npos && start_pos < Restore_List.size()) {
 			restore_path = Restore_List.substr(start_pos, end_pos - start_pos);
-			restore_part = Find_Partition_By_Path(restore_path);
-			if (restore_part != NULL) {
-				if (restore_part->Mount_Read_Only) {
-					gui_msg(Msg(msg::kError, "restore_read_only=Cannot restore {1} -- mounted read only.")(restore_part->Backup_Display_Name));
+			part_settings.Part = Find_Partition_By_Path(restore_path);
+			if (part_settings.Part != NULL) {
+				part_settings.Backup_FileName = part_settings.Part->Backup_Name + "." +  part_settings.Part->Current_File_System + ".win";
+				if (part_settings.Part->Mount_Read_Only) {
+					gui_msg(Msg(msg::kError, "restore_read_only=Cannot restore {1} -- mounted read only.")(part_settings.Part->Backup_Display_Name));
 					return false;
 				}
-				if (check_md5 > 0 && !restore_part->Check_MD5(Restore_Name))
+
+				if (check_md5 > 0 && !part_settings.Part->Check_MD5(Restore_Name))
 					return false;
-				partition_count++;
-				total_restore_size += restore_part->Get_Restore_Size(Restore_Name);
-				if (restore_part->Has_SubPartition) {
+				part_settings.partition_count++;
+				part_settings.total_restore_size += part_settings.Part->Get_Restore_Size(&part_settings);
+				if (part_settings.Part->Has_SubPartition) {
 					std::vector<TWPartition*>::iterator subpart;
 
 					for (subpart = Partitions.begin(); subpart != Partitions.end(); subpart++) {
-						if ((*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == restore_part->Mount_Point) {
+						if ((*subpart)->Is_SubPartition && (*subpart)->SubPartition_Of == part_settings.Part->Mount_Point) {
 							if (check_md5 > 0 && !(*subpart)->Check_MD5(Restore_Name))
 								return false;
-							total_restore_size += (*subpart)->Get_Restore_Size(Restore_Name);
+							part_settings.total_restore_size += (*subpart)->Get_Restore_Size(&part_settings);
 						}
 					}
 				}
@@ -909,25 +976,29 @@ int TWPartitionManager::Run_Restore(const string& Restore_Name) {
 		}
 	}
 
-	if (partition_count == 0) {
+	if (part_settings.partition_count == 0) {
 		gui_err("no_part_restore=No partitions selected for restore.");
 		return false;
 	}
 
-	gui_msg(Msg("restore_part_count=Restoring {1} partitions...")(partition_count));
-	gui_msg(Msg("total_restore_size=Total restore size is {1}MB")(total_restore_size / 1048576));
+	gui_msg(Msg("restore_part_count=Restoring {1} partitions...")(part_settings.partition_count));
+	gui_msg(Msg("total_restore_size=Total restore size is {1}MB")(part_settings.total_restore_size / 1048576));
 	DataManager::SetProgress(0.0);
-	ProgressTracking progress(total_restore_size);
+	ProgressTracking progress(part_settings.total_restore_size);
+	part_settings.progress = &progress;
 
 	start_pos = 0;
 	if (!Restore_List.empty()) {
 		end_pos = Restore_List.find(";", start_pos);
 		while (end_pos != string::npos && start_pos < Restore_List.size()) {
 			restore_path = Restore_List.substr(start_pos, end_pos - start_pos);
-			restore_part = Find_Partition_By_Path(restore_path);
-			if (restore_part != NULL) {
-				partition_count++;
-				if (!Restore_Partition(restore_part, Restore_Name, &progress))
+
+			part_settings.Part = Find_Partition_By_Path(restore_path);
+			if (part_settings.Part != NULL) {
+				part_settings.partition_count++;
+				part_settings.Backup_FileName = part_settings.Part->Backup_Name + "." +  part_settings.Part->Current_File_System + ".win";
+				part_settings.Full_Backup_Path = part_settings.Backup_Folder + "/" + part_settings.Backup_FileName + "/";
+				if (!Restore_Partition(&part_settings))
 					return false;
 			} else {
 				gui_msg(Msg(msg::kError, "restore_unable_locate=Unable to locate '{1}' partition for restoring.")(restore_path));
@@ -942,6 +1013,7 @@ int TWPartitionManager::Run_Restore(const string& Restore_Name) {
 	time(&rStop);
 	gui_msg(Msg(msg::kHighlight, "restore_completed=[RESTORE COMPLETED IN {1} SECONDS]")((int)difftime(rStop,rStart)));
 	DataManager::SetValue("tw_file_progress", "");
+
 	return true;
 }
 
@@ -2170,21 +2242,23 @@ bool TWPartitionManager::Remove_MTP_Storage(unsigned int Storage_ID) {
 	return false;
 }
 
-bool TWPartitionManager::Flash_Image(string Filename) {
+bool TWPartitionManager::Flash_Image(PartitionSettings *part_settings) {
 	int check, partition_count = 0;
 	TWPartition* flash_part = NULL;
-	string Flash_List, flash_path;
+	string Flash_List, flash_path, full_filename;
 	size_t start_pos = 0, end_pos = 0;
 
+	full_filename = part_settings->Restore_Name + "/" + part_settings->Backup_FileName;
+
 	gui_msg("image_flash_start=[IMAGE FLASH STARTED]");
-	gui_msg(Msg("img_to_flash=Image to flash: '{1}'")(Filename));
+	gui_msg(Msg("img_to_flash=Image to flash: '{1}'")(full_filename));
 
-	if (!TWFunc::Path_Exists(Filename)) {
-		if (!Mount_By_Path(Filename, true)) {
+	if (!TWFunc::Path_Exists(full_filename)) {
+		if (!Mount_By_Path(full_filename, true)) {
 			return false;
 		}
-		if (!TWFunc::Path_Exists(Filename)) {
-			gui_msg(Msg(msg::kError, "unable_to_locate=Unable to locate {1}.")(Filename));
+		if (!TWFunc::Path_Exists(full_filename)) {
+			gui_msg(Msg(msg::kError, "unable_to_locate=Unable to locate {1}.")(full_filename));
 			return false;
 		}
 	}
@@ -2218,7 +2292,7 @@ bool TWPartitionManager::Flash_Image(string Filename) {
 
 	DataManager::SetProgress(0.0);
 	if (flash_part) {
-		if (!flash_part->Flash_Image(Filename))
+		if (!flash_part->Flash_Image(part_settings))
 			return false;
 	} else {
 		gui_err("invalid_flash=Invalid flash partition specified.");