1 files changed, 187 insertions, 12 deletions

diff --git a/src/restore.c b/src/restore.c
index fd6fec2..ce9d41b 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -24,6 +24,7 @@
 #include <string.h>
 #include <libimobiledevice/restore.h>
 
+#include "tss.h"
 #include "restore.h"
 #include "idevicerestore.h"
 
@@ -44,6 +45,8 @@
 #define WAIT_FOR_DEVICE        33
 #define LOAD_NOR               36
 
+static int restore_device_connected = 0;
+
 int restore_check_mode(const char* uuid) {
 	char* type = NULL;
 	uint64_t version = 0;
@@ -77,8 +80,81 @@ int restore_check_mode(const char* uuid) {
 	return 0;
 }
 
+void restore_device_callback(const idevice_event_t* event, void* user_data) {
+	if (event->event == IDEVICE_DEVICE_ADD) {
+		restore_device_connected = 1;
+
+	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
+		restore_device_connected = 0;
+	}
+}
+
+int restore_open_with_timeout(const char* uuid, idevice_t* device, restored_client_t* restore) {
+	int i = 0;
+	int attempt = 10;
+	char* type = NULL;
+	uint64_t version = 0;
+	idevice_t context = NULL;
+	restored_client_t client = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	restored_error_t restore_error = RESTORE_E_SUCCESS;
+
+	*device = NULL;
+	*restore = NULL;
+
+	device_error = idevice_event_subscribe(&restore_device_callback, NULL);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		error("ERROR: Unable to subscribe to device events\n");
+		return -1;
+	}
+
+	for (i = 1; i <= attempt; i++) {
+		if (restore_device_connected == 1) {
+			break;
+		}
+
+		if (i == attempt) {
+			error("ERROR: Unable to connect to device in restore mode\n");
+		}
+
+		sleep(2);
+	}
+
+	device_error = idevice_new(&context, uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	restore_error = restored_client_new(context, &client, "idevicerestore");
+	if (restore_error != RESTORE_E_SUCCESS) {
+		idevice_event_unsubscribe();
+		idevice_free(context);
+		return -1;
+	}
+
+	restore_error = restored_query_type(client, &type, &version);
+	if (restore_error != RESTORE_E_SUCCESS) {
+		restored_client_free(client);
+		idevice_event_unsubscribe();
+		idevice_free(context);
+		return -1;
+	}
+
+	*device = context;
+	*restore = client;
+	return 0;
+}
+
+void restore_close(idevice_t device, restored_client_t restore) {
+	if (restore)
+		restored_client_free(restore);
+	if (device)
+		idevice_free(device);
+	idevice_event_unsubscribe();
+}
+
 const char* restore_progress_string(unsigned int operation) {
-	switch(operation) {
+	switch (operation) {
 	case CREATE_PARTITION_MAP:
 		return "Creating partition map";
 
@@ -126,7 +202,6 @@ const char* restore_progress_string(unsigned int operation) {
 	}
 }
 
-
 int restore_handle_progress_msg(restored_client_t client, plist_t msg) {
 	plist_t node = NULL;
 	uint64_t operation = 0;
@@ -163,7 +238,7 @@ int restore_handle_status_msg(restored_client_t client, plist_t msg) {
 	return 0;
 }
 
-int restore_send_filesystem(idevice_t device, restored_client_t client, const char *filesystem) {
+int restore_send_filesystem(idevice_t device, restored_client_t client, const char* filesystem) {
 	int i = 0;
 	char buffer[0x1000];
 	uint32_t recv_bytes = 0;
@@ -327,7 +402,7 @@ int restore_send_filesystem(idevice_t device, restored_client_t client, const ch
 	return ret;
 }
 
-int restore_send_kernelcache(restored_client_t client, char *kernel_data, int len) {
+int restore_send_kernelcache(restored_client_t client, char* kernel_data, int len) {
 	info("Sending kernelcache\n");
 
 	plist_t kernelcache_node = plist_new_data(kernel_data, len);
@@ -347,19 +422,17 @@ int restore_send_kernelcache(restored_client_t client, char *kernel_data, int le
 	return 0;
 }
 
-int restore_send_nor(restored_client_t client, char* ipsw, plist_t tss) {
+int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss) {
 	char* llb_path = NULL;
-	char* llb_blob = NULL;
-	if (get_tss_data_by_name(tss, "LLB", &llb_path, &llb_blob) < 0) {
+	if (tss_get_entry_path(tss, "LLB", &llb_path) < 0) {
 		error("ERROR: Unable to get LLB info from TSS response\n");
 		return -1;
 	}
 
 	char* llb_filename = strstr(llb_path, "LLB");
 	if (llb_filename == NULL) {
-		error("ERROR: Unable to extrac firmware path from LLB filename\n");
+		error("ERROR: Unable to extract firmware path from LLB filename\n");
 		free(llb_path);
-		free(llb_blob);
 		return -1;
 	}
 
@@ -378,7 +451,6 @@ int restore_send_nor(restored_client_t client, char* ipsw, plist_t tss) {
 	if (ipsw_extract_to_memory(ipsw, manifest_file, &manifest_data, &manifest_size) < 0) {
 		error("ERROR: Unable to extract firmware manifest from ipsw\n");
 		free(llb_path);
-		free(llb_blob);
 		return -1;
 	}
 
@@ -392,7 +464,7 @@ int restore_send_nor(restored_client_t client, char* ipsw, plist_t tss) {
 	if (filename != NULL) {
 		memset(firmware_filename, '\0', sizeof(firmware_filename));
 		snprintf(firmware_filename, sizeof(firmware_filename), "%s/%s", firmware_path, filename);
-		if (get_signed_component_by_path(ipsw, tss, firmware_filename, &llb_data, &llb_size) < 0) {
+		if (get_signed_component(ipsw, tss, firmware_filename, &llb_data, &llb_size) < 0) {
 			error("ERROR: Unable to get signed LLB\n");
 			return -1;
 		}
@@ -407,7 +479,7 @@ int restore_send_nor(restored_client_t client, char* ipsw, plist_t tss) {
 	while (filename != NULL) {
 		memset(firmware_filename, '\0', sizeof(firmware_filename));
 		snprintf(firmware_filename, sizeof(firmware_filename), "%s/%s", firmware_path, filename);
-		if (get_signed_component_by_path(ipsw, tss, firmware_filename, &nor_data, &nor_size) < 0) {
+		if (get_signed_component(ipsw, tss, firmware_filename, &nor_data, &nor_size) < 0) {
 			error("ERROR: Unable to get signed firmware %s\n", firmware_filename);
 			break;
 		}
@@ -436,3 +508,106 @@ int restore_send_nor(restored_client_t client, char* ipsw, plist_t tss) {
 	plist_free(dict);
 	return 0;
 }
+
+int restore_device(const char* uuid, const char* ipsw, plist_t tss, const char* filesystem) {
+	idevice_t device = NULL;
+	idevice_error_t device_error = idevice_new(&device, uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		error("ERROR: Unable to open device\n");
+		plist_free(tss);
+		return -1;
+	}
+
+	restored_client_t restore = NULL;
+	restored_error_t restore_error = restored_client_new(device, &restore, "idevicerestore");
+	if (restore_error != RESTORE_E_SUCCESS) {
+		error("ERROR: Unable to start restored client\n");
+		plist_free(tss);
+		idevice_free(device);
+		return -1;
+	}
+
+	char* type = NULL;
+	uint64_t version = 0;
+	if (restored_query_type(restore, &type, &version) != RESTORE_E_SUCCESS) {
+		error("ERROR: Device is not in restore mode. QueryType returned \"%s\"\n", type);
+		plist_free(tss);
+		restored_client_free(restore);
+		idevice_free(device);
+		return -1;
+	}
+	info("Device has successfully entered restore mode\n");
+
+	/* start restore process */
+	char* kernelcache = NULL;
+	info("Restore protocol version is %llu.\n", version);
+	restore_error = restored_start_restore(restore);
+	if (restore_error == RESTORE_E_SUCCESS) {
+		while (!idevicerestore_quit) {
+			plist_t message = NULL;
+			restore_error = restored_receive(restore, &message);
+			plist_t msgtype_node = plist_dict_get_item(message, "MsgType");
+			if (msgtype_node && PLIST_STRING == plist_get_node_type(msgtype_node)) {
+				char *msgtype = NULL;
+				plist_get_string_val(msgtype_node, &msgtype);
+				if (!strcmp(msgtype, "ProgressMsg")) {
+					restore_error = restore_handle_progress_msg(restore, message);
+
+				} else if (!strcmp(msgtype, "DataRequestMsg")) {
+					// device is requesting data to be sent
+					plist_t datatype_node = plist_dict_get_item(message, "DataType");
+					if (datatype_node && PLIST_STRING == plist_get_node_type(datatype_node)) {
+						char *datatype = NULL;
+						plist_get_string_val(datatype_node, &datatype);
+						if (!strcmp(datatype, "SystemImageData")) {
+							restore_send_filesystem(device, restore, filesystem);
+
+						} else if (!strcmp(datatype, "KernelCache")) {
+							int kernelcache_size = 0;
+							char* kernelcache_data = NULL;
+							char* kernelcache_path = NULL;
+							if (tss_get_entry_path(tss, "KernelCache", &kernelcache_path) < 0) {
+								error("ERROR: Unable to find kernelcache path\n");
+								return -1;
+							}
+
+							if (get_signed_component(ipsw, tss, kernelcache_path, &kernelcache_data, &kernelcache_size) < 0) {
+								error("ERROR: Unable to get kernelcache file\n");
+								return -1;
+							}
+							restore_send_kernelcache(restore, kernelcache_data, kernelcache_size);
+							free(kernelcache_data);
+
+						} else if (!strcmp(datatype, "NORData")) {
+							restore_send_nor(restore, ipsw, tss);
+
+						} else {
+							// Unknown DataType!!
+							error("Unknown DataType\n");
+							return -1;
+						}
+					}
+
+				} else if (!strcmp(msgtype, "StatusMsg")) {
+					restore_error = restore_handle_status_msg(restore, message);
+
+				} else {
+					info("Received unknown message type: %s\n", msgtype);
+				}
+			}
+
+			if (RESTORE_E_SUCCESS != restore_error) {
+				error("Invalid return status %d\n", restore_error);
+				//idevicerestore_quit = 1;
+			}
+
+			plist_free(message);
+		}
+	} else {
+		error("ERROR: Could not start restore. %d\n", restore_error);
+	}
+
+	restored_client_free(restore);
+	idevice_free(device);
+	return 0;
+}