Merge branch 'martin'

6 files changed, 186 insertions, 61 deletions

diff --git a/src/asr.c b/src/asr.c
index 2331660..f48170a 100644
--- a/src/asr.c
+++ b/src/asr.c
@@ -87,7 +87,9 @@ int asr_receive(idevice_connection_t asr, plist_t* data) {
 
 	*data = request;
 
-	debug("Received %d bytes:\n%s\n", size, buffer);
+	debug("Received %d bytes:\n", size);
+	if (idevicerestore_debug)
+		debug_plist(request);
 	free(buffer);
 	return 0;
 }
@@ -104,7 +106,8 @@ int asr_send(idevice_connection_t asr, plist_t* data) {
 	}
 
 	debug("Sent %d bytes:\n", size);
-	debug_plist(data);
+	if (idevicerestore_debug)
+		debug_plist(*data);
 	free(buffer);
 	return 0;
 }
@@ -119,6 +122,8 @@ int asr_send_buffer(idevice_connection_t asr, const char* data, uint32_t size) {
 		return -1;
 	}
 
+	debug("Sent %d bytes buffer\n", bytes);
+
 	return 0;
 }
 
@@ -137,6 +142,7 @@ int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
 	plist_t packet = NULL;
 	plist_t packet_info = NULL;
 	plist_t payload_info = NULL;
+	int attempts = 0;
 
 	file = fopen(filesystem, "rb");
 	if (file == NULL) {
@@ -172,6 +178,17 @@ int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
 			return -1;
 		}
 
+		if (packet == NULL) {
+			if (attempts < 5) {
+				info("Retrying to receive validation packet... %d\n", attempts);
+				attempts++;
+				sleep(1);
+				continue;
+			}
+		}
+
+		attempts = 0;
+
 		node = plist_dict_get_item(packet, "Command");
 		if (!node || plist_get_node_type(node) != PLIST_STRING) {
 			error("ERROR: Unable to find command node in validation request\n");
@@ -181,10 +198,7 @@ int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
 
 		if (!strcmp(command, "OOBData")) {
 			asr_handle_oob_data_request(asr, packet, file);
-
-
 			plist_free(packet);
-
 		} else if(!strcmp(command, "Payload")) {
 			plist_free(packet);
 			break;
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index 7982ed7..00c35c1 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -81,20 +81,20 @@ int main(int argc, char* argv[]) {
 			return 0;
 
 		case 'd':
-			client->flags &= FLAG_DEBUG;
+			client->flags |= FLAG_DEBUG;
 			idevicerestore_debug = 1;
 			break;
 
 		case 'e':
-			client->flags &= FLAG_ERASE;
+			client->flags |= FLAG_ERASE;
 			break;
 
 		case 'c':
-			client->flags &= FLAG_CUSTOM;
+			client->flags |= FLAG_CUSTOM;
 			break;
 
 		case 'x':
-			client->flags &= FLAG_EXCLUDE;
+			client->flags |= FLAG_EXCLUDE;
 			break;
 
 		case 'u':
@@ -149,15 +149,18 @@ int main(int argc, char* argv[]) {
 	// extract buildmanifest
 	plist_t buildmanifest = NULL;
 	info("Extracting BuildManifest from IPSW\n");
-	if (extract_buildmanifest(client, ipsw, &buildmanifest) < 0) {
+	if (ipsw_extract_build_manifest(ipsw, &buildmanifest) < 0) {
 		error("ERROR: Unable to extract BuildManifest from %s\n", ipsw);
 		return -1;
 	}
 
+	/* print iOS information from the manifest */
+	build_manifest_print_information(buildmanifest);
+
 	// devices are listed in order from oldest to newest
 	// so we'll need their ECID
 	if (client->device->index > DEVICE_IPOD2G) {
-		debug("Creating TSS request\n");
+		debug("Getting device's ECID for TSS request\n");
 		// fetch the device's ECID for the TSS request
 		if (get_ecid(client, &client->ecid) < 0) {
 			error("ERROR: Unable to find device ECID\n");
@@ -170,7 +173,7 @@ int main(int argc, char* argv[]) {
 	client->tss = NULL;
 	plist_t build_identity = NULL;
 	if (client->flags & FLAG_ERASE) {
-		build_identity = get_build_identity(client, buildmanifest, 0);
+		build_identity = build_manifest_get_build_identity(buildmanifest, 0);
 		if (build_identity == NULL) {
 			error("ERROR: Unable to find any build identities\n");
 			plist_free(buildmanifest);
@@ -181,13 +184,16 @@ int main(int argc, char* argv[]) {
 		// and list the valid ones
 		int i = 0;
 		int valid_builds = 0;
-		int build_count = get_build_count(buildmanifest);
+		int build_count = build_manifest_get_identity_count(buildmanifest);
 		for (i = 0; i < build_count; i++) {
-			build_identity = get_build_identity(client, buildmanifest, i);
+			build_identity = build_manifest_get_build_identity(buildmanifest, i);
 			valid_builds++;
 		}
 	}
 
+	/* print information about current build identity */
+	build_identity_print_information(build_identity);
+
 	if (client->flags & FLAG_CUSTOM > 0) {
 		if (client->device->index > DEVICE_IPOD2G) {
 			if (get_shsh_blobs(client, ecid, build_identity, &client->tss) < 0) {
@@ -206,7 +212,7 @@ int main(int argc, char* argv[]) {
 
 	// Extract filesystem from IPSW and return its name
 	char* filesystem = NULL;
-	if (extract_filesystem(client, client->ipsw, build_identity, &filesystem) < 0) {
+	if (ipsw_extract_filesystem(client->ipsw, build_identity, &filesystem) < 0) {
 		error("ERROR: Unable to extract filesystem from IPSW\n");
 		if (client->tss)
 			plist_free(client->tss);
@@ -455,32 +461,9 @@ int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
 	return 0;
 }
 
-int extract_buildmanifest(struct idevicerestore_client_t* client, const char* ipsw, plist_t* buildmanifest) {
-	int size = 0;
-	char* data = NULL;
-	int device = client->device->index;
-
-	/* older devices don't require personalized firmwares and use a BuildManifesto.plist */
-	if (ipsw_extract_to_memory(ipsw, "BuildManifesto.plist", &data, &size) == 0) {
-		plist_from_xml(data, size, buildmanifest);
-		return 0;
-	}
-
-	data = NULL;
-	size = 0;
-
-	/* whereas newer devices do not require personalized firmwares and use a BuildManifest.plist */
-	if (ipsw_extract_to_memory(ipsw, "BuildManifest.plist", &data, &size) == 0) {
-		plist_from_xml(data, size, buildmanifest);
-		return 0;
-	}
-
-	return -1;
-}
-
-plist_t get_build_identity(struct idevicerestore_client_t* client, plist_t buildmanifest, uint32_t identity) {
+plist_t build_manifest_get_build_identity(plist_t build_manifest, uint32_t identity) {
 	// fetch build identities array from BuildManifest
-	plist_t build_identities_array = plist_dict_get_item(buildmanifest, "BuildIdentities");
+	plist_t build_identities_array = plist_dict_get_item(build_manifest, "BuildIdentities");
 	if (!build_identities_array || plist_get_node_type(build_identities_array) != PLIST_ARRAY) {
 		error("ERROR: Unable to find build identities node\n");
 		return NULL;
@@ -523,9 +506,9 @@ int get_shsh_blobs(struct idevicerestore_client_t* client, uint64_t ecid, plist_
 	return 0;
 }
 
-int get_build_count(plist_t buildmanifest) {
+int build_manifest_get_identity_count(plist_t build_manifest) {
 	// fetch build identities array from BuildManifest
-	plist_t build_identities_array = plist_dict_get_item(buildmanifest, "BuildIdentities");
+	plist_t build_identities_array = plist_dict_get_item(build_manifest, "BuildIdentities");
 	if (!build_identities_array || plist_get_node_type(build_identities_array) != PLIST_ARRAY) {
 		error("ERROR: Unable to find build identities node\n");
 		return -1;
@@ -535,7 +518,7 @@ int get_build_count(plist_t buildmanifest) {
 	return plist_array_get_size(build_identities_array);
 }
 
-int extract_filesystem(struct idevicerestore_client_t* client, const char* ipsw, plist_t build_identity, char** filesystem) {
+int ipsw_extract_filesystem(const char* ipsw, plist_t build_identity, char** filesystem) {
 	char* filename = NULL;
 
 	if (build_identity_get_component_path(build_identity, "OS", &filename) < 0) {
@@ -616,6 +599,73 @@ int ipsw_get_component_by_path(const char* ipsw, plist_t tss, const char* path,
 	return 0;
 }
 
+void build_manifest_print_information(plist_t build_manifest) {
+	char* value = NULL;
+	plist_t node = NULL;
+
+	node = plist_dict_get_item(build_manifest, "ProductVersion");
+	if (!node || plist_get_node_type(node) != PLIST_STRING) {
+		error("ERROR: Unable to find ProductVersion node\n");
+		return;
+	}
+	plist_get_string_val(node, &value);
+
+	info("Product Version: %s\n", value);
+	free(value);
+
+	node = plist_dict_get_item(build_manifest, "ProductBuildVersion");
+	if (!node || plist_get_node_type(node) != PLIST_STRING) {
+		error("ERROR: Unable to find ProductBuildVersion node\n");
+		return;
+	}
+	plist_get_string_val(node, &value);
+
+	info("Product Build: %s\n", value);
+	free(value);
+
+	node = NULL;
+}
+
+void build_identity_print_information(plist_t build_identity) {
+	char* value = NULL;
+	plist_t info_node = NULL;
+	plist_t node = NULL;
+
+	info_node = plist_dict_get_item(build_identity, "Info");
+	if (!info_node || plist_get_node_type(info_node) != PLIST_DICT) {
+		error("ERROR: Unable to find Info node\n");
+		return;
+	}
+
+	node = plist_dict_get_item(info_node, "Variant");
+	if (!node || plist_get_node_type(node) != PLIST_STRING) {
+		error("ERROR: Unable to find Variant node\n");
+		return;
+	}
+	plist_get_string_val(node, &value);
+
+	info("Variant: %s\n", value);
+	free(value);
+
+	node = plist_dict_get_item(info_node, "RestoreBehavior");
+	if (!node || plist_get_node_type(node) != PLIST_STRING) {
+		error("ERROR: Unable to find RestoreBehavior node\n");
+		return;
+	}
+	plist_get_string_val(node, &value);
+
+	if (!strcmp(value, "Erase"))
+		info("This restore will erase your device data.\n");
+
+	if (!strcmp(value, "Update"))
+		info("This restore will update your device without loosing data.\n");
+
+	free(value);
+
+	info_node = NULL;
+	node = NULL;
+}
+
 int build_identity_get_component_path(plist_t build_identity, const char* component, char** path) {
 	char* filename = NULL;
 
diff --git a/src/idevicerestore.h b/src/idevicerestore.h
index f529b5b..4f59a02 100644
--- a/src/idevicerestore.h
+++ b/src/idevicerestore.h
@@ -34,16 +34,17 @@ extern "C" {
 void usage(int argc, char* argv[]);
 int check_mode(struct idevicerestore_client_t* client);
 int check_device(struct idevicerestore_client_t* client);
-int get_build_count(plist_t buildmanifest);
 int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid);
 int get_bdid(struct idevicerestore_client_t* client, uint32_t* bdid);
 int get_cpid(struct idevicerestore_client_t* client, uint32_t* cpid);
-int extract_buildmanifest(struct idevicerestore_client_t* client, const char* ipsw, plist_t* buildmanifest);
-plist_t get_build_identity(struct idevicerestore_client_t* client, plist_t buildmanifest, uint32_t identity);
 int get_shsh_blobs(struct idevicerestore_client_t* client, uint64_t ecid, plist_t build_identity, plist_t* tss);
-int extract_filesystem(struct idevicerestore_client_t* client, const char* ipsw, plist_t buildmanifest, char** filesystem);
-int ipsw_get_component_by_path(const char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size);
+void build_manifest_print_information(plist_t build_manifest);
+plist_t build_manifest_get_build_identity(plist_t build_manifest, uint32_t identity);
+int build_manifest_get_build_count(plist_t build_manifest);
+void build_identity_print_information(plist_t build_identity);
 int build_identity_get_component_path(plist_t build_identity, const char* component, char** path);
+int ipsw_extract_filesystem(const char* ipsw, plist_t build_identity, char** filesystem);
+int ipsw_get_component_by_path(const char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size);
 
 #ifdef __cplusplus
 }
diff --git a/src/ipsw.c b/src/ipsw.c
index f08e2fd..9cd7290 100644
--- a/src/ipsw.c
+++ b/src/ipsw.c
@@ -173,6 +173,28 @@ int ipsw_extract_to_memory(const char* ipsw, const char* infile, char** pbuffer,
 	return 0;
 }
 
+int ipsw_extract_build_manifest(const char* ipsw, plist_t* buildmanifest) {
+	int size = 0;
+	char* data = NULL;
+
+	/* older devices don't require personalized firmwares and use a BuildManifesto.plist */
+	if (ipsw_extract_to_memory(ipsw, "BuildManifesto.plist", &data, &size) == 0) {
+		plist_from_xml(data, size, buildmanifest);
+		return 0;
+	}
+
+	data = NULL;
+	size = 0;
+
+	/* whereas newer devices do not require personalized firmwares and use a BuildManifest.plist */
+	if (ipsw_extract_to_memory(ipsw, "BuildManifest.plist", &data, &size) == 0) {
+		plist_from_xml(data, size, buildmanifest);
+		return 0;
+	}
+
+	return -1;
+}
+
 void ipsw_close(ipsw_archive* archive) {
 	if (archive != NULL) {
 		zip_unchange_all(archive->zip);
diff --git a/src/ipsw.h b/src/ipsw.h
index cd11406..f1694ef 100644
--- a/src/ipsw.h
+++ b/src/ipsw.h
@@ -28,6 +28,7 @@ extern "C" {
 
 #include <zip.h>
 #include <stdint.h>
+#include <plist/plist.h>
 
 typedef struct {
 	int index;
@@ -37,6 +38,7 @@ typedef struct {
 } ipsw_file;
 
 int ipsw_extract_to_memory(const char* ipsw, const char* infile, char** pbuffer, uint32_t* psize);
+int ipsw_extract_build_manifest(const char* ipsw, plist_t* buildmanifest);
 void ipsw_free_file(ipsw_file* file);
 
 #ifdef __cplusplus
diff --git a/src/restore.c b/src/restore.c
index 5d45296..bb3ed96 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -29,6 +29,7 @@
 #include "common.h"
 #include "restore.h"
 
+#define WAIT_FOR_STORAGE       11
 #define CREATE_PARTITION_MAP   12
 #define CREATE_FILESYSTEM      13
 #define RESTORE_IMAGE          14
@@ -39,6 +40,8 @@
 #define UPDATE_BASEBAND        20
 #define FINIALIZE_NAND         21
 #define MODIFY_BOOTARGS        26
+#define LOAD_KERNEL_CACHE      27
+#define PARTITION_NAND_DEVICE  28
 #define WAIT_FOR_NAND          29
 #define UNMOUNT_FILESYSTEM     30
 #define WAIT_FOR_DEVICE        33
@@ -178,7 +181,7 @@ void restore_device_callback(const idevice_event_t* event, void* userdata) {
 
 	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
 		restore_device_connected = 0;
-		client->flags &= FLAG_QUIT;
+		client->flags |= FLAG_QUIT;
 	}
 }
 
@@ -225,7 +228,7 @@ int restore_open_with_timeout(struct idevicerestore_client_t* client) {
 		}
 	}
 
-	device_error = idevice_event_subscribe(&restore_device_callback, NULL);
+	device_error = idevice_event_subscribe(&restore_device_callback, client);
 	if (device_error != IDEVICE_E_SUCCESS) {
 		error("ERROR: Unable to subscribe to device events\n");
 		return -1;
@@ -271,6 +274,9 @@ int restore_open_with_timeout(struct idevicerestore_client_t* client) {
 
 const char* restore_progress_string(unsigned int operation) {
 	switch (operation) {
+	case WAIT_FOR_STORAGE:
+		return "Waiting for Storage Device...";
+
 	case CREATE_PARTITION_MAP:
 		return "Creating partition map";
 
@@ -304,12 +310,18 @@ const char* restore_progress_string(unsigned int operation) {
 	case UNMOUNT_FILESYSTEM:
 		return "Unmounting filesystems";
 
+	case PARTITION_NAND_DEVICE:
+		return "Partition NAND device";
+
 	case WAIT_FOR_NAND:
 		return "Waiting for NAND...";
 
 	case WAIT_FOR_DEVICE:
 		return "Waiting for Device...";
 
+	case LOAD_KERNEL_CACHE:
+		return "Loading kernelcache...";
+
 	case LOAD_NOR:
 		return "Loading NOR data to flash";
 
@@ -339,7 +351,6 @@ int restore_handle_progress_msg(restored_client_t client, plist_t msg) {
 
 	if ((progress > 0) && (progress < 100)) {
 		print_progress_bar((double) progress);
-
 	} else {
 		info("%s\n", restore_progress_string(operation));
 	}
@@ -348,8 +359,27 @@ int restore_handle_progress_msg(restored_client_t client, plist_t msg) {
 }
 
 int restore_handle_status_msg(restored_client_t client, plist_t msg) {
+	uint64_t value = 0;
 	info("Got status message\n");
 	debug_plist(msg);
+
+	plist_t node = plist_dict_get_item(msg, "Status");
+	plist_get_uint_val(node, &value);
+
+	switch(value) {
+		case 0:
+			info("Status: Restore Finished\n");
+			break;
+		case 6:
+			info("Status: Disk Failure\n");
+			break;
+		case 14:
+			info("Status: Fail\n");
+			break;
+		default:
+			info("Unknown status message.\n");
+	}
+
 	return 0;
 }
 
@@ -366,8 +396,7 @@ int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	}
 	info("Connected to ASR\n");
 
-	// we don't really need to do anything with this,
-	// we're just clearing the output buffer
+	/* receive Initiate command message */
 	if (asr_receive(asr, &data) < 0) {
 		error("ERROR: Unable to receive data from ASR\n");
 		asr_close(asr);
@@ -530,11 +559,12 @@ int restore_send_nor(restored_client_t restore, struct idevicerestore_client_t*
 	}
 	plist_dict_insert_item(dict, "NorImageData", norimage_array);
 
-	debug_plist(dict);
+	if (idevicerestore_debug)
+		debug_plist(dict);
 
 	ret = restored_send(restore, dict);
 	if (ret != RESTORE_E_SUCCESS) {
-		error("ERROR: Unable to send kernelcache data\n");
+		error("ERROR: Unable to send NOR image data data\n");
 		plist_free(dict);
 		return -1;
 	}
@@ -569,19 +599,22 @@ int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idev
 		}
 
 		else if (!strcmp(type, "NORData")) {
-			if(client->flags & FLAG_EXCLUDE > 0) {
+			if((client->flags & FLAG_EXCLUDE) == 0) {
+				info("Sending NORData\n");
 				if(restore_send_nor(restore, client, build_identity) < 0) {
 					error("ERROR: Unable to send NOR data\n");
 					return -1;
 				}
 			} else {
-				client->flags &= 1;
+				info("Not sending NORData... Quitting...\n");
+				client->flags |= FLAG_QUIT;
 			}
 
 		} else {
 			// Unknown DataType!!
 			debug("Unknown data request received\n");
-			debug_plist(message);
+			if (idevicerestore_debug)
+				debug_plist(message);
 		}
 	}
 	return 0;
@@ -606,6 +639,7 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 	info("Device has successfully entered restore mode\n");
 
 	restore = client->restore->client;
+	device = client->restore->device;
 
 	// start the restore process
 	restore_error = restored_start_restore(restore);
@@ -629,7 +663,8 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 		node = plist_dict_get_item(message, "MsgType");
 		if (!node || plist_get_node_type(node) != PLIST_STRING) {
 			debug("Unknown message received\n");
-			debug_plist(message);
+			if (idevicerestore_debug)
+				debug_plist(message);
 			plist_free(message);
 			message = NULL;
 			continue;
@@ -659,14 +694,15 @@ int restore_device(struct idevicerestore_client_t* client, plist_t build_identit
 		// at least the "previous error logs" messages usually end up here
 		else {
 			debug("Unknown message type received\n");
-			debug_plist(message);
+			if (idevicerestore_debug)
+				debug_plist(message);
 		}
 
 		// finally, if any of these message handlers returned -1 then we encountered
 		// an unrecoverable error, so we need to bail.
 		if (error < 0) {
 			error("ERROR: Unable to successfully restore device\n");
-			client->flags &= FLAG_QUIT;
+			client->flags |= FLAG_QUIT;
 		}
 
 		plist_free(message);