5 files changed, 86 insertions, 33 deletions

diff --git a/gui/action.cpp b/gui/action.cpp
index 4d9c9df59..dd38b3175 100644
--- a/gui/action.cpp
+++ b/gui/action.cpp
@@ -670,7 +670,13 @@ int GUIAction::doAction(Action action, int isThreaded /* = 0 */)
 				if (simulate) {
 					simulate_progress_bar();
 				} else {
-					system("injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash");
+					TWPartition* Boot = PartitionManager.Find_Partition_By_Path("/boot");
+					if (Boot == NULL || Boot->Current_File_System != "emmc")
+						system("injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash");
+					else {
+						string injectcmd = "injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash bd=" + Boot->Actual_Block_Device;
+						system(injectcmd.c_str());
+					}
 					ui_print("TWRP injection complete.\n");
 				}
 			}
@@ -1042,6 +1048,22 @@ int GUIAction::doAction(Action action, int isThreaded /* = 0 */)
 					ret = 1; // failure
 				else if (wipe_cache)
 					PartitionManager.Wipe_By_Path("/cache");
+				if (DataManager::GetIntValue(TW_HAS_INJECTTWRP) == 1 && DataManager::GetIntValue(TW_INJECT_AFTER_ZIP) == 1) {
+					operation_start("ReinjectTWRP");
+					ui_print("Injecting TWRP into boot image...\n");
+					if (simulate) {
+						simulate_progress_bar();
+					} else {
+						TWPartition* Boot = PartitionManager.Find_Partition_By_Path("/boot");
+						if (Boot == NULL || Boot->Current_File_System != "emmc")
+							system("injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash");
+						else {
+							string injectcmd = "injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash bd=" + Boot->Actual_Block_Device;
+							system(injectcmd.c_str());
+						}
+						ui_print("TWRP injection complete.\n");
+					}
+				}
 			}
 			operation_end(ret, simulate);
 			return 0;
diff --git a/injecttwrp/Android.mk b/injecttwrp/Android.mk
index 9c39c297c..2557523a2 100644
--- a/injecttwrp/Android.mk
+++ b/injecttwrp/Android.mk
@@ -1,13 +1,13 @@
-LOCAL_PATH:= $(call my-dir)

-include $(CLEAR_VARS)

-

-ifeq ($(TW_INCLUDE_INJECTTWRP), true)

-	LOCAL_SRC_FILES:= \

-		injecttwrp.c

-	LOCAL_CFLAGS:= -g -c -W

-	LOCAL_MODULE:=injecttwrp

-	LOCAL_MODULE_TAGS:= eng

-	LOCAL_MODULE_CLASS := RECOVERY_EXECUTABLES

-	LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/sbin

-	include $(BUILD_EXECUTABLE)

-endif

+LOCAL_PATH:= $(call my-dir)
+include $(CLEAR_VARS)
+
+ifeq ($(TW_INCLUDE_INJECTTWRP), true)
+	LOCAL_SRC_FILES:= \
+		injecttwrp.c
+	LOCAL_CFLAGS:= -g -c -W
+	LOCAL_MODULE:=injecttwrp
+	LOCAL_MODULE_TAGS:= eng
+	LOCAL_MODULE_CLASS := RECOVERY_EXECUTABLES
+	LOCAL_MODULE_PATH := $(TARGET_RECOVERY_ROOT_OUT)/sbin
+	include $(BUILD_EXECUTABLE)
+endif
diff --git a/injecttwrp/injecttwrp.c b/injecttwrp/injecttwrp.c
index 0daaf5696..731f9190f 100644
--- a/injecttwrp/injecttwrp.c
+++ b/injecttwrp/injecttwrp.c
@@ -297,10 +297,10 @@ int find_gzip_recovery_ramdisk(char *boot_image, unsigned long *ramdisk_address)
 }
 
 int main(int argc, char** argv) {
-	int arg_error = 0, delete_ind = 0, return_val;
+	int arg_error = 0, delete_ind = 0, return_val, index, len;
 	unsigned long address2;
 	unsigned char regular_check[8] = "ANDROID!";
-	char boot_image[512], backup_image[512];
+	char boot_image[512], backup_image[512], boot_block_device[512], command[512];
 
 	printf("-- InjectTWRP Recovery Ramdisk Injection Tool for Samsung devices. --\n");
 	printf("--                  by Dees_Troy and Team Win                      --\n");
@@ -308,17 +308,26 @@ int main(int argc, char** argv) {
 	printf("--                Bringing some win to Samsung!                    --\n");
 	printf("--        This tool comes with no warranties whatsoever!           --\n");
 	printf("--        Use at your own risk and always keep a backup!           --\n\n");
-	printf("Version 0.1 beta\n\n");
+	printf("Version 0.2 beta\n\n");
 
 	// Parse the arguments
-	if (argc < 2 || argc > 5)
+	if (argc < 2 || argc > 6)
 		arg_error = 1;
 	else {
-		if ((argc == 2 || argc == 3) && (strcmp(argv[1], "-b") == 0 || strcmp(argv[1], "--backup") == 0)) {
+		strcpy(boot_block_device, "boot");
+		for (index = 1; index < argc; index++) {
+			len = strlen(argv[index]);
+			if (len > 3 && strncmp(argv[index], "bd=", 3) == 0) {
+				strcpy(boot_block_device, argv[index] + 3);
+				index = argc;
+			}
+		}
+		if ((argc >= 2 && argc <= 4) && (strcmp(argv[1], "-b") == 0 || strcmp(argv[1], "--backup") == 0)) {
 			// Backup existing boot image
 			printf("Dumping boot image...\n");
 #ifdef INJECT_USE_TMP
-			system("dump_image boot /tmp/original_boot.img");
+			sprintf(command, "dump_image %s /tmp/original_boot.img", boot_block_device);
+			system(command);
 			strcpy(boot_image, "/tmp/original_boot.img");
 
 			if (argc == 2)
@@ -327,7 +336,8 @@ int main(int argc, char** argv) {
 				strcpy(backup_image, argv[2]);
 #else
 			system("mount /cache");
-			system("dump_image boot /cache/original_boot.img");
+			sprintf(command, "dump_image %s /cache/original_boot.img", boot_block_device);
+			system(command);
 			strcpy(boot_image, "/cache/original_boot.img");
 
 			if (argc == 2)
@@ -356,14 +366,16 @@ int main(int argc, char** argv) {
 			if (strcmp(argv[1], "-d") == 0 || strcmp(argv[1], "--dump") == 0) {
 				printf("Dumping boot image...\n");
 #ifdef INJECT_USE_TMP
-				system("dump_image boot /tmp/original_boot.img");
+				sprintf(command, "dump_image %s /tmp/original_boot.img", boot_block_device);
+				system(command);
 				strcpy(boot_image, "/tmp/original_boot.img");
 #else
 				system("mount /cache");
-				system("dump_image boot /cache/original_boot.img");
+				sprintf(command, "dump_image %s /cache/original_boot.img", boot_block_device);
+				system(command);
 				strcpy(boot_image, "/cache/original_boot.img");
-#endif
 				delete_ind = -1;
+#endif
 			} else
 				strcpy(boot_image, argv[1]);
 
@@ -389,13 +401,11 @@ int main(int argc, char** argv) {
 				system("rm /cache/original_boot.img");
 			}
 
-			if (argc == 5 && (strcmp(argv[4], "-f") == 0 || strcmp(argv[4], "--flash") == 0)) {
-				char command[512];
-
+			if (argc >= 5 && (strcmp(argv[4], "-f") == 0 || strcmp(argv[4], "--flash") == 0)) {
 				printf("Flashing new image...\n");
-				system("erase_image boot"); // Needed because flash_image checks the header and the header sometimes is the same while the ramdisks are different
-				strcpy(command, "flash_image boot ");
-				strcat(command, argv[3]);
+				sprintf(command, "erase_image %s", boot_block_device);
+				system(command);
+				sprintf(command, "flash_image %s %s", boot_block_device, argv[3]);
 				system(command);
 				printf("Flash complete.\n");
 			}
@@ -413,7 +423,9 @@ int main(int argc, char** argv) {
 		printf("injecttwrp --dump ramdisk-recovery.img outputboot.img [--flash]\n");
 		printf("--dump will use dump_image to dump your existing boot image\n");
 		printf("--flash will use flash_image to flash the new boot image\n\n");
-		printf("NOTE: dump_image, erase_image, and flash_image must already be installed!\n");
+		printf("NOTE: dump_image, erase_image, and flash_image must already be installed!\n\n");
+		printf("If needed you can add bd=/dev/block/mmcblk0p5 to indicate the location\n");
+		printf("of the boot partition on emmc devices as the final parameter.\n");
 		return 0;
 	}
 
diff --git a/openrecoveryscript.cpp b/openrecoveryscript.cpp
index 9317ae7bc..47bf5d6dd 100644
--- a/openrecoveryscript.cpp
+++ b/openrecoveryscript.cpp
@@ -72,7 +72,7 @@ int OpenRecoveryScript::check_for_script_file(void) {
 
 int OpenRecoveryScript::run_script_file(void) {
 	FILE *fp = fopen(SCRIPT_FILE_TMP, "r");
-	int ret_val = 0, cindex, line_len, i, remove_nl;
+	int ret_val = 0, cindex, line_len, i, remove_nl, install_cmd = 0;
 	char script_line[SCRIPT_COMMAND_SIZE], command[SCRIPT_COMMAND_SIZE],
 		 value[SCRIPT_COMMAND_SIZE], mount[SCRIPT_COMMAND_SIZE],
 		 value1[SCRIPT_COMMAND_SIZE], value2[SCRIPT_COMMAND_SIZE];
@@ -112,6 +112,7 @@ int OpenRecoveryScript::run_script_file(void) {
 			if (strcmp(command, "install") == 0) {
 				// Install Zip
 				ret_val = Install_Command(value);
+				install_cmd = -1;
 			} else if (strcmp(command, "wipe") == 0) {
 				// Wipe
 				if (strcmp(value, "cache") == 0 || strcmp(value, "/cache") == 0) {
@@ -345,6 +346,17 @@ int OpenRecoveryScript::run_script_file(void) {
 		LOGE("Error opening script file '%s'\n", SCRIPT_FILE_TMP);
 		return 1;
 	}
+	if (install_cmd && DataManager_GetIntValue(TW_HAS_INJECTTWRP) == 1 && DataManager_GetIntValue(TW_INJECT_AFTER_ZIP) == 1) {
+		ui_print("Injecting TWRP into boot image...\n");
+		TWPartition* Boot = PartitionManager.Find_Partition_By_Path("/boot");
+		if (Boot == NULL || Boot->Current_File_System != "emmc")
+			system("injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash");
+		else {
+			string injectcmd = "injecttwrp --dump /tmp/backup_recovery_ramdisk.img /tmp/injected_boot.img --flash bd=" + Boot->Actual_Block_Device;
+			system(injectcmd.c_str());
+		}
+		ui_print("TWRP injection complete.\n");
+	}
 	return ret_val;
 }
 
diff --git a/recovery.cpp b/recovery.cpp
index b8fc45f70..b1333eaa5 100644
--- a/recovery.cpp
+++ b/recovery.cpp
@@ -893,8 +893,15 @@ main(int argc, char **argv) {
 
 #ifdef TW_INCLUDE_INJECTTWRP
 	// Back up TWRP Ramdisk if needed:
+	TWPartition* Boot = PartitionManager.Find_Partition_By_Path("/boot");
+
 	LOGI("Backing up TWRP ramdisk...\n");
-	system("injecttwrp --backup /tmp/backup_recovery_ramdisk.img");
+	if (Boot == NULL || Boot->Current_File_System != "emmc")
+		system("injecttwrp --backup /tmp/backup_recovery_ramdisk.img");
+	else {
+		string injectcmd = "injecttwrp --backup /tmp/backup_recovery_ramdisk.img bd=" + Boot->Actual_Block_Device;
+		system(injectcmd.c_str());
+	}
 	LOGI("Backup of TWRP ramdisk done.\n");
 #endif