libidevicerecovery: implemented progress callback logic

9 files changed, 108 insertions, 12 deletions

diff --git a/src/asr.c b/src/asr.c
index 401bc51..96d7396 100644
--- a/src/asr.c
+++ b/src/asr.c
@@ -27,6 +27,7 @@
 
 #include "asr.h"
 #include "idevicerestore.h"
+#include "common.h"
 
 #define ASR_VERSION 1
 #define ASR_STREAM_ID 1
@@ -103,6 +104,15 @@ int asr_open_with_timeout(idevice_t device, asr_client_t* asr) {
 	return 0;
 }
 
+void asr_set_progress_callback(asr_client_t asr, asr_progress_cb_t cbfunc, void* userdata)
+{
+	if (!asr) {
+		return;
+	}
+	asr->progress_cb = cbfunc;
+	asr->progress_cb_data = userdata;
+}
+
 int asr_receive(asr_client_t asr, plist_t* data) {
 	uint32_t size = 0;
 	char* buffer = NULL;
@@ -387,8 +397,11 @@ int asr_send_payload(asr_client_t asr, const char* filesystem) {
 		}
 
 		bytes += size;
-		progress = ((double) bytes/ (double) length) * 100.0;
-		print_progress_bar(progress);
+		progress = ((double) bytes/ (double) length);
+		if (asr->progress_cb && ((int)(progress*100) > asr->lastprogress)) {
+			asr->progress_cb(progress, asr->progress_cb_data);
+			asr->lastprogress = (int)(progress*100);
+		}
 	}
 
 	// if last chunk wasn't terminated with a checksum we do it here
diff --git a/src/asr.h b/src/asr.h
index f4752b1..2524d48 100644
--- a/src/asr.h
+++ b/src/asr.h
@@ -28,13 +28,19 @@ extern "C" {
 
 #include <libimobiledevice/libimobiledevice.h>
 
+typedef void (*asr_progress_cb_t)(double, void*);
+
 struct asr_client {
 	idevice_connection_t connection;
 	uint8_t checksum_chunks;
+	int lastprogress;
+	asr_progress_cb_t progress_cb;
+	void* progress_cb_data;
 };
 typedef struct asr_client *asr_client_t;
 
 int asr_open_with_timeout(idevice_t device, asr_client_t* asr);
+void asr_set_progress_callback(asr_client_t asr, asr_progress_cb_t, void* userdata);
 int asr_send(asr_client_t asr, plist_t* data);
 int asr_receive(asr_client_t asr, plist_t* data);
 int asr_send_buffer(asr_client_t asr, const char* data, uint32_t size);
diff --git a/src/common.c b/src/common.c
index a2ff7b3..2b364cb 100644
--- a/src/common.c
+++ b/src/common.c
@@ -158,3 +158,15 @@ int mkdir_with_parents(const char *dir, int mode)
 	}
 	return res;
 }
+
+void idevicerestore_progress(struct idevicerestore_client_t* client, int step, double progress)
+{
+	if(client && client->progress_cb) {
+		client->progress_cb(step, progress, client->progress_cb_data);
+	} else {
+		// we don't want to be too verbose in regular idevicerestore.
+		if ((step == RESTORE_STEP_UPLOAD_FS) || (step == RESTORE_STEP_FLASH_FS) || (step == RESTORE_STEP_FLASH_NOR)) {
+			print_progress_bar(100.0f * progress);
+		}
+	}
+}
diff --git a/src/common.h b/src/common.h
index 1c3e2fa..dae1eea 100644
--- a/src/common.h
+++ b/src/common.h
@@ -29,6 +29,8 @@ extern "C" {
 #include <plist/plist.h>
 #include <libirecovery.h>
 
+#include "idevicerestore.h"
+
 #define info(...) printf(__VA_ARGS__)
 #define error(...) fprintf(stderr, __VA_ARGS__)
 #define debug(...) if(idevicerestore_debug) fprintf(stderr, __VA_ARGS__)
@@ -85,6 +87,8 @@ struct idevicerestore_client_t {
 	char* build;
 	char* restore_boot_args;
 	char* cache_dir;
+	idevicerestore_progress_cb_t progress_cb;
+	void* progress_cb_data;
 };
 
 static struct idevicerestore_mode_t idevicerestore_modes[] = {
@@ -119,6 +123,8 @@ char *generate_guid();
 
 int mkdir_with_parents(const char *dir, int mode);
 
+void idevicerestore_progress(struct idevicerestore_client_t* client, int step, double progress);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index acdd572..3b95702 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -177,6 +177,8 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		irecv_set_debug_level(1);
 	}
 
+	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.0);
+
 	// update version data (from cache, or apple if too old)
 	load_version_data(client);
 
@@ -185,6 +187,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
 		return -1;
 	}
+	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.1);
 	info("Found device in %s mode\n", client->mode->string);
 
 	if (client->mode->index == MODE_WTF) {
@@ -258,6 +261,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		error("ERROR: Unable to discover device type\n");
 		return -1;
 	}
+	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.2);
 	info("Identified device as %s\n", client->device->product);
 
 	if ((client->flags & FLAG_PWN) && (client->mode->index != MODE_DFU)) {
@@ -297,6 +301,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			client->ipsw = ipsw;
 		}
 	}
+	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.6);
 
 	if (client->flags & FLAG_NOACTION) {
 		return 0;
@@ -334,6 +339,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			return -1;
 		}
 	}
+	idevicerestore_progress(client, RESTORE_STEP_DETECT, 0.8);
 
 	/* check if device type is supported by the given build manifest */
 	if (build_manifest_check_compatibility(buildmanifest, client->device->product) < 0) {
@@ -518,6 +524,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	/* print information about current build identity */
 	build_identity_print_information(build_identity);
 
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.0);
 	/* retrieve shsh blobs if required */
 	if (tss_enabled) {
 		debug("Getting device's ECID for TSS request\n");
@@ -587,6 +594,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		/* fix empty dicts */
 		fixup_tss(client->tss);
 	}
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.1);
 
 	// if the device is in normal mode, place device into recovery mode
 	if (client->mode->index == MODE_NORMAL) {
@@ -694,6 +702,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		}
 	}
 
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.3);
 
 	// if the device is in DFU mode, place device into recovery mode
 	if (client->mode->index == MODE_DFU) {
@@ -744,6 +753,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		/* FIXME: Probably better to detect if the device is back then */
 		sleep(7);
 	}
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.5);
 
 	if (client->build[0] > '8') {
 		// we need another tss request with nonce.
@@ -794,6 +804,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			fixup_tss(client->tss);
 		}
 	}
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.7);
 
 	// now finally do the magic to put the device into restore mode
 	if (client->mode->index == MODE_RECOVERY) {
@@ -813,6 +824,7 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 			return -2;
 		}
 	}
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.9);
 
 	// device is finally in restore mode, let's do this
 	if (client->mode->index == MODE_RESTORE) {
@@ -831,6 +843,11 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 		unlink(filesystem);
 
 	info("DONE\n");
+
+	if (result == 0) {
+		idevicerestore_progress(client, RESTORE_NUM_STEPS-1, 1.0);
+	}
+
 	return result;
 }
 
@@ -924,6 +941,14 @@ void idevicerestore_set_ipsw(struct idevicerestore_client_t* client, const char*
 	}
 }
 
+void idevicerestore_set_progress_callback(struct idevicerestore_client_t* client, idevicerestore_progress_cb_t cbfunc, void* userdata)
+{
+	if (!client)
+		return;
+	client->progress_cb = cbfunc;
+	client->progress_cb_data = userdata;
+}
+
 #ifndef IDEVICERESTORE_NOMAIN
 int main(int argc, char* argv[]) {
 	int opt = 0;
diff --git a/src/idevicerestore.h b/src/idevicerestore.h
index 3201280..c7c80dc 100644
--- a/src/idevicerestore.h
+++ b/src/idevicerestore.h
@@ -29,8 +29,6 @@ extern "C" {
 #include <stdint.h>
 #include <plist/plist.h>
 
-#include "common.h"
-
 // the flag with value 1 is reserved for internal use only. don't use it.
 #define FLAG_DEBUG           1 << 1
 #define FLAG_ERASE           1 << 2
@@ -43,6 +41,17 @@ extern "C" {
 
 struct idevicerestore_client_t;
 
+enum {
+	RESTORE_STEP_DETECT = 0,
+	RESTORE_STEP_PREPARE,
+	RESTORE_STEP_UPLOAD_FS,
+	RESTORE_STEP_FLASH_FS,
+	RESTORE_STEP_FLASH_NOR,
+	RESTORE_NUM_STEPS
+};
+
+typedef void (*idevicerestore_progress_cb_t)(int step, double step_progress, void* userdata);
+
 struct idevicerestore_client_t* idevicerestore_client_new();
 void idevicerestore_client_free(struct idevicerestore_client_t* client);
 
@@ -51,6 +60,8 @@ void idevicerestore_set_udid(struct idevicerestore_client_t* client, const char*
 void idevicerestore_set_flags(struct idevicerestore_client_t* client, int flags);
 void idevicerestore_set_ipsw(struct idevicerestore_client_t* client, const char* path);
 
+void idevicerestore_set_progress_callback(struct idevicerestore_client_t* client, idevicerestore_progress_cb_t cbfunc, void* userdata);
+
 int idevicerestore_start(struct idevicerestore_client_t* client);
 
 void usage(int argc, char* argv[]);
diff --git a/src/ipsw.c b/src/ipsw.c
index 84547b3..249800d 100644
--- a/src/ipsw.c
+++ b/src/ipsw.c
@@ -27,6 +27,7 @@
 
 #include "ipsw.h"
 #include "locking.h"
+#include "common.h"
 #include "idevicerestore.h"
 
 #define BUFSIZE 0x100000
diff --git a/src/restore.c b/src/restore.c
index 0d8f266..2529afe 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -31,6 +31,7 @@
 #include "mbn.h"
 #include "tss.h"
 #include "restore.h"
+#include "common.h"
 
 #define WAIT_FOR_STORAGE       11
 #define CREATE_PARTITION_MAP   12
@@ -527,7 +528,7 @@ int restore_handle_previous_restore_log_msg(restored_client_t client, plist_t ms
 	return 0;
 }
 
-int restore_handle_progress_msg(restored_client_t client, plist_t msg) {
+int restore_handle_progress_msg(struct idevicerestore_client_t* client, plist_t msg) {
 	plist_t node = NULL;
 	uint64_t progress = 0;
 	uint64_t operation = 0;
@@ -546,11 +547,21 @@ int restore_handle_progress_msg(restored_client_t client, plist_t msg) {
 	}
 	plist_get_uint_val(node, &progress);
 
-	if ((progress > 0) && (progress < 100)) {
+	if ((progress > 0) && (progress <= 100)) {
 		if (operation != lastop) {
 			info("%s (%d)\n", restore_progress_string(operation), (int)operation);
 		}
-		print_progress_bar((double) progress);
+		switch ((int)operation) {
+		case 14:
+			idevicerestore_progress(client, RESTORE_STEP_FLASH_FS, progress / 100.0);
+			break;
+		case 18:
+			idevicerestore_progress(client, RESTORE_STEP_FLASH_NOR, progress / 100.0);
+			break;
+		default:
+			debug("Unhandled progress operation %d\n", (int)operation);
+			break;
+		}
 	} else {
 		info("%s (%d)\n", restore_progress_string(operation), (int)operation);
 	}
@@ -654,7 +665,15 @@ int restore_handle_bb_update_status_msg(restored_client_t client, plist_t msg)
 	return result;
 }
 
-int restore_send_filesystem(idevice_t device, const char* filesystem) {
+void restore_asr_progress_cb(double progress, void* userdata)
+{
+	struct idevicerestore_client_t* client = (struct idevicerestore_client_t*)userdata;
+	if (client) {
+		idevicerestore_progress(client, RESTORE_STEP_UPLOAD_FS, progress);
+	}
+}
+
+int restore_send_filesystem(struct idevicerestore_client_t* client, idevice_t device, const char* filesystem) {
 	int i = 0;
 	FILE* file = NULL;
 	plist_t data = NULL;
@@ -669,6 +688,8 @@ int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	}
 	info("Connected to ASR\n");
 
+	asr_set_progress_callback(asr, restore_asr_progress_cb, (void*)client);
+
 	// this step sends requested chunks of data from various offsets to asr so
 	// it can validate the filesystem before installing it
 	info("Validating the filesystem\n");
@@ -1406,7 +1427,7 @@ int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idev
 
 		// this request is sent when restored is ready to receive the filesystem
 		if (!strcmp(type, "SystemImageData")) {
-			if(restore_send_filesystem(device, filesystem) < 0) {
+			if(restore_send_filesystem(client, device, filesystem) < 0) {
 				error("ERROR: Unable to send filesystem\n");
 				return -2;
 			}
@@ -1606,6 +1627,7 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 		return -1;
 	}
 	plist_free(opts);
+	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 1.0);
 
 	// this is the restore process loop, it reads each message in from
 	// restored and passes that data on to it's specific handler
@@ -1651,7 +1673,7 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 		// progress notification messages sent by the restored inform the client
 		// of it's current operation and sometimes percent of progress is complete
 		else if (!strcmp(type, "ProgressMsg")) {
-			error = restore_handle_progress_msg(restore, message);
+			error = restore_handle_progress_msg(client, message);
 		}
 
 		// status messages usually indicate the current state of the restored
diff --git a/src/restore.h b/src/restore.h
index 071fa8f..2389b3b 100644
--- a/src/restore.h
+++ b/src/restore.h
@@ -48,14 +48,14 @@ void restore_client_free(struct idevicerestore_client_t* client);
 int restore_reboot(struct idevicerestore_client_t* client);
 const char* restore_progress_string(unsigned int operation);
 int restore_handle_status_msg(restored_client_t client, plist_t msg);
-int restore_handle_progress_msg(restored_client_t client, plist_t msg);
+int restore_handle_progress_msg(struct idevicerestore_client_t* client, plist_t msg);
 int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idevice_t device, restored_client_t restore, plist_t message, plist_t build_identity, const char* filesystem);
 int restore_send_nor(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity);
 int restore_send_root_ticket(restored_client_t restore, struct idevicerestore_client_t* client);
 int restore_send_kernelcache(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity);
 int restore_device(struct idevicerestore_client_t* client, plist_t build_identity, const char* filesystem);
 int restore_open_with_timeout(struct idevicerestore_client_t* client);
-int restore_send_filesystem(idevice_t device, const char* filesystem);
+int restore_send_filesystem(struct idevicerestore_client_t* client, idevice_t device, const char* filesystem);
 
 
 #ifdef __cplusplus