1 files changed, 176 insertions, 14 deletions

diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index 34ba6aea6..79068f79d 100755
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -602,16 +602,15 @@ TWPartition* TWPartitionManager::Find_Partition_By_Block_Device(const string& Bl
 	return NULL;
 }
 
-int TWPartitionManager::Check_Backup_Name(bool Display_Error) {
+int TWPartitionManager::Check_Backup_Name(const std::string& Backup_Name, bool Display_Error, bool Must_Be_Unique) {
 	// Check the backup name to ensure that it is the correct size and contains only valid characters
 	// and that a backup with that name doesn't already exist
 	char backup_name[MAX_BACKUP_NAME_LEN];
 	char backup_loc[255], tw_image_dir[255];
 	int copy_size;
 	int index, cur_char;
-	string Backup_Name, Backup_Loc;
+	string Backup_Loc;
 
-	DataManager::GetValue(TW_BACKUP_NAME, Backup_Name);
 	copy_size = Backup_Name.size();
 	// Check size
 	if (copy_size > MAX_BACKUP_NAME_LEN) {
@@ -640,17 +639,20 @@ int TWPartitionManager::Check_Backup_Name(bool Display_Error) {
 		}
 	}
 
-	// Check to make sure that a backup with this name doesn't already exist
-	DataManager::GetValue(TW_BACKUPS_FOLDER_VAR, Backup_Loc);
-	strcpy(backup_loc, Backup_Loc.c_str());
-	sprintf(tw_image_dir,"%s/%s", backup_loc, Backup_Name.c_str());
-	if (TWFunc::Path_Exists(tw_image_dir)) {
-		if (Display_Error)
-			gui_err("backup_name_exists=A backup with that name already exists!");
+	if (Must_Be_Unique) {
+		// Check to make sure that a backup with this name doesn't already exist
+		DataManager::GetValue(TW_BACKUPS_FOLDER_VAR, Backup_Loc);
+		strcpy(backup_loc, Backup_Loc.c_str());
+		sprintf(tw_image_dir,"%s/%s", backup_loc, Backup_Name.c_str());
+		if (TWFunc::Path_Exists(tw_image_dir)) {
+			if (Display_Error)
+				gui_err("backup_name_exists=A backup with that name already exists!");
 
-		return -4;
+			return -4;
+		}
+		// Backup is unique
 	}
-	// No problems found, return 0
+	// No problems found
 	return 0;
 }
 
@@ -2210,6 +2212,22 @@ void TWPartitionManager::Get_Partition_List(string ListType, std::vector<Partiti
 				Partition_List->push_back(part);
 			}
 		}
+		if (DataManager::GetIntValue("tw_has_repack_tools") != 0 && DataManager::GetIntValue("tw_has_boot_slots") != 0) {
+			TWPartition* boot = Find_Partition_By_Path("/boot");
+			if (boot) {
+				// Allow flashing kernels and ramdisks
+				struct PartitionList repack_ramdisk;
+				repack_ramdisk.Display_Name = gui_lookup("install_twrp_ramdisk", "Install Recovery Ramdisk");
+				repack_ramdisk.Mount_Point = "/repack_ramdisk";
+				repack_ramdisk.selected = 0;
+				Partition_List->push_back(repack_ramdisk);
+				/*struct PartitionList repack_kernel; For now let's leave repacking kernels under advanced only
+				repack_kernel.Display_Name = gui_lookup("install_kernel", "Install Kernel");
+				repack_kernel.Mount_Point = "/repack_kernel";
+				repack_kernel.selected = 0;
+				Partition_List->push_back(repack_kernel);*/
+			}
+		}
 	} else {
 		LOGERR("Unknown list type '%s' requested for TWPartitionManager::Get_Partition_List\n", ListType.c_str());
 	}
@@ -2512,6 +2530,21 @@ bool TWPartitionManager::Flash_Image(string& path, string& filename) {
 		}
 	}
 
+	DataManager::GetValue("tw_flash_partition", Flash_List);
+	Repack_Type repack = REPLACE_NONE;
+	if (Flash_List == "/repack_ramdisk;") {
+		repack = REPLACE_RAMDISK;
+	} else if (Flash_List == "/repack_kernel;") {
+		repack = REPLACE_KERNEL;
+	}
+	if (repack != REPLACE_NONE) {
+		Repack_Options_struct Repack_Options;
+		Repack_Options.Type = repack;
+		Repack_Options.Disable_Verity = false;
+		Repack_Options.Disable_Force_Encrypt = false;
+		Repack_Options.Backup_First = DataManager::GetIntValue("tw_repack_backup_first") != 0;
+		return Repack_Images(full_filename, Repack_Options);
+	}
 	PartitionSettings part_settings;
 	part_settings.Backup_Folder = path;
 	unsigned long long total_bytes = TWFunc::Get_File_Size(full_filename);
@@ -2519,9 +2552,7 @@ bool TWPartitionManager::Flash_Image(string& path, string& filename) {
 	part_settings.progress = &progress;
 	part_settings.adbbackup = false;
 	part_settings.PM_Method = PM_RESTORE;
-
 	gui_msg("calc_restore=Calculating restore details...");
-	DataManager::GetValue("tw_flash_partition", Flash_List);
 	if (!Flash_List.empty()) {
 		end_pos = Flash_List.find(";", start_pos);
 		while (end_pos != string::npos && start_pos < Flash_List.size()) {
@@ -2995,3 +3026,134 @@ void TWPartitionManager::Coldboot() {
 	if (sysfs_entries.size() > 0)
 		Coldboot_Scan(&sysfs_entries, "/sys/block", 0);
 }
+
+bool TWPartitionManager::Prepare_Empty_Folder(const std::string& Folder) {
+	if (TWFunc::Path_Exists(Folder))
+		TWFunc::removeDir(Folder, false);
+	return TWFunc::Recursive_Mkdir(Folder);
+}
+
+bool TWPartitionManager::Prepare_Repack(TWPartition* Part, const std::string& Temp_Folder_Destination, const bool Create_Backup, const std::string& Backup_Name) {
+	if (!Part) {
+		LOGERR("Partition was null!\n");
+		return false;
+	}
+	if (!Prepare_Empty_Folder(Temp_Folder_Destination))
+		return false;
+	std::string target_image = Temp_Folder_Destination + "boot.img";
+	PartitionSettings part_settings;
+	part_settings.Part = Part;
+	if (Create_Backup) {
+		if (Check_Backup_Name(Backup_Name, true, false) != 0)
+			return false;
+		DataManager::GetValue(TW_BACKUPS_FOLDER_VAR, part_settings.Backup_Folder);
+		part_settings.Backup_Folder = part_settings.Backup_Folder + "/" + TWFunc::Get_Current_Date() + " " + Backup_Name + "/";
+		if (!TWFunc::Recursive_Mkdir(part_settings.Backup_Folder))
+			return false;
+	} else
+		part_settings.Backup_Folder = Temp_Folder_Destination;
+	part_settings.adbbackup = false;
+	part_settings.generate_digest = false;
+	part_settings.generate_md5 = false;
+	part_settings.PM_Method = PM_BACKUP;
+	part_settings.progress = NULL;
+	pid_t not_a_pid = 0;
+	if (!Part->Backup(&part_settings, &not_a_pid))
+		return false;
+	std::string backed_up_image = part_settings.Backup_Folder;
+	backed_up_image += Part->Backup_FileName;
+	target_image = Temp_Folder_Destination + "boot.img";
+	if (Create_Backup) {
+		std::string source = part_settings.Backup_Folder + Part->Backup_FileName;
+		if (TWFunc::copy_file(source, target_image, 0644) != 0) {
+			LOGERR("Failed to copy backup file '%s' to temp folder target '%s'\n", source.c_str(), target_image.c_str());
+			return false;
+		}
+	} else {
+		if (rename(backed_up_image.c_str(), target_image.c_str()) != 0) {
+			LOGERR("Failed to rename '%s' to '%s'\n", backed_up_image.c_str(), target_image.c_str());
+			return false;
+		}
+	}
+	return Prepare_Repack(target_image, Temp_Folder_Destination, false, false);
+}
+
+bool TWPartitionManager::Prepare_Repack(const std::string& Source_Path, const std::string& Temp_Folder_Destination, const bool Copy_Source, const bool Create_Destination) {
+	if (Create_Destination) {
+		if (!Prepare_Empty_Folder(Temp_Folder_Destination))
+			return false;
+	}
+	if (Copy_Source) {
+		std::string destination = Temp_Folder_Destination + "/boot.img";
+		if (TWFunc::copy_file(Source_Path, destination, 0644))
+			return false;
+	}
+	std::string command = "cd " + Temp_Folder_Destination + " && /sbin/magiskboot --unpack -h " + Source_Path;
+	if (TWFunc::Exec_Cmd(command) != 0) {
+		LOGINFO("Error unpacking %s!\n", Source_Path.c_str());
+		gui_msg(Msg(msg::kError, "unpack_error=Error unpacking image."));
+		return false;
+	}
+	return true;
+}
+
+bool TWPartitionManager::Repack_Images(const std::string& Target_Image, const struct Repack_Options_struct& Repack_Options) {
+	if (!TWFunc::Path_Exists("/sbin/magiskboot")) {
+		LOGERR("Image repacking tool not present in this TWRP build!");
+		return false;
+	}
+	DataManager::SetProgress(0);
+	TWPartition* part = PartitionManager.Find_Partition_By_Path("/boot");
+	if (part)
+		gui_msg(Msg("unpacking_image=Unpacking {1}...")(part->Display_Name));
+	else {
+		gui_msg(Msg(msg::kError, "unable_to_locate=Unable to locate {1}.")("/boot"));
+		return false;
+	}
+	if (!PartitionManager.Prepare_Repack(part, REPACK_ORIG_DIR, Repack_Options.Backup_First, gui_lookup("repack", "Repack")))
+		return false;
+	DataManager::SetProgress(.25);
+	gui_msg(Msg("unpacking_image=Unpacking {1}...")(Target_Image));
+	if (!PartitionManager.Prepare_Repack(Target_Image, REPACK_NEW_DIR, true))
+		return false;
+	DataManager::SetProgress(.5);
+	gui_msg(Msg("repacking_image=Repacking {1}...")(part->Display_Name));
+	std::string path = REPACK_NEW_DIR;
+	if (Repack_Options.Type == REPLACE_KERNEL) {
+		// When we replace the kernel, what we really do is copy the boot partition ramdisk into the new image's folder
+		if (TWFunc::copy_file(REPACK_ORIG_DIR "ramdisk.cpio", REPACK_NEW_DIR "ramdisk.cpio", 0644)) {
+			LOGERR("Failed to copy ramdisk\n");
+			return false;
+		}
+	} else if (Repack_Options.Type == REPLACE_RAMDISK) {
+		// Repack the ramdisk
+		if (TWFunc::copy_file(REPACK_NEW_DIR "ramdisk.cpio", REPACK_ORIG_DIR "ramdisk.cpio", 0644)) {
+			LOGERR("Failed to copy ramdisk\n");
+			return false;
+		}
+		path = REPACK_ORIG_DIR;
+	} else {
+		LOGERR("Invalid repacking options specified\n");
+		return false;
+	}
+	if (Repack_Options.Disable_Verity)
+		LOGERR("Disabling verity is not implemented yet\n");
+	if (Repack_Options.Disable_Force_Encrypt)
+		LOGERR("Disabling force encrypt is not implemented yet\n");
+	std::string command = "cd " + path + " && /sbin/magiskboot --repack " + path + "boot.img";
+	if (TWFunc::Exec_Cmd(command) != 0) {
+		gui_msg(Msg(msg::kError, "repack_error=Error repacking image."));
+		return false;
+	}
+	DataManager::SetProgress(.75);
+	std::string file = "new-boot.img";
+	DataManager::SetValue("tw_flash_partition", "/boot;");
+	if (!PartitionManager.Flash_Image(path, file)) {
+		LOGINFO("Error flashing new image\n");
+		return false;
+	}
+	DataManager::SetProgress(1);
+	TWFunc::removeDir(REPACK_ORIG_DIR, false);
+	TWFunc::removeDir(REPACK_NEW_DIR, false);
+	return true;
+}