Allow restores that start in Restore mode. Allow specification of custom APTickets.

5 files changed, 128 insertions, 44 deletions

diff --git a/src/common.h b/src/common.h
index e8e4d44..d943568 100644
--- a/src/common.h
+++ b/src/common.h
@@ -103,6 +103,8 @@ struct idevicerestore_client_t {
 	int build_major;
 	char* restore_boot_args;
 	char* cache_dir;
+	unsigned char* root_ticket;
+	int root_ticket_len;
 	idevicerestore_progress_cb_t progress_cb;
 	void* progress_cb_data;
 	irecv_device_event_context_t irecv_e_ctx;
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index 0ee2bde..5601082 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -73,6 +73,8 @@ static struct option longopts[] = {
 	{ "no-action", no_argument,     NULL, 'n' },
 	{ "cache-path", required_argument, NULL, 'C' },
 	{ "no-input", no_argument,      NULL, 'y' },
+	{ "restore-mode", no_argument,  NULL, 'R' },
+	{ "ticket", required_argument,  NULL, 'T' },
 	{ NULL, 0, NULL, 0 }
 };
 
@@ -117,6 +119,8 @@ static void usage(int argc, char* argv[], int err)
 	" -t, --shsh       Fetch TSS record and save to .shsh file, then exit\n" \
 	" -k, --keep-pers  Write personalized components to files for debugging\n" \
 	" -p, --pwn        Put device in pwned DFU mode and exit (limera1n devices only)\n" \
+	" -R, --restore-mode  Allow restoring from Restore mode\n" \
+	" -T, --ticket PATH   Use file at PATH to send as AP ticket\n" \
 	"\n" \
 	"Homepage: <" PACKAGE_URL ">\n",
 	(name ? name + 1 : argv[0]));
@@ -557,21 +561,29 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	}
 
 	if (client->mode->index == MODE_RESTORE) {
-		if (restore_reboot(client) < 0) {
-			error("ERROR: Unable to exit restore mode\n");
-			return -2;
-		}
+		if (client->flags & FLAG_ALLOW_RESTORE_MODE) {
+			tss_enabled = 0;
+			if (!client->root_ticket) {
+				client->root_ticket = (void*)strdup("");
+				client->root_ticket_len = 0;
+			}
+		} else {
+			if (restore_reboot(client) < 0) {
+				error("ERROR: Unable to exit restore mode\n");
+				return -2;
+			}
 
-		// we need to refresh the current mode again
-		mutex_lock(&client->device_event_mutex);
-		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 60000);
-		if (client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+			// we need to refresh the current mode again
+			mutex_lock(&client->device_event_mutex);
+			cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 60000);
+			if (client->mode == &idevicerestore_modes[MODE_UNKNOWN] || (client->flags & FLAG_QUIT)) {
+				mutex_unlock(&client->device_event_mutex);
+				error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
+				return -1;
+			}
+			info("Found device in %s mode\n", client->mode->string);
 			mutex_unlock(&client->device_event_mutex);
-			error("ERROR: Unable to discover device mode. Please make sure a device is attached.\n");
-			return -1;
 		}
-		info("Found device in %s mode\n", client->mode->string);
-		mutex_unlock(&client->device_event_mutex);
 	}
 
 	// verify if ipsw file exists
@@ -1302,17 +1314,19 @@ int idevicerestore_start(struct idevicerestore_client_t* client)
 	}
 	idevicerestore_progress(client, RESTORE_STEP_PREPARE, 0.9);
 
-	mutex_lock(&client->device_event_mutex);
-	info("Waiting for device to enter restore mode...\n");
-	cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 180000);
-	if (client->mode != &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT)) {
+	if (client->mode->index != MODE_RESTORE) {
+		mutex_lock(&client->device_event_mutex);
+		info("Waiting for device to enter restore mode...\n");
+		cond_wait_timeout(&client->device_event_cond, &client->device_event_mutex, 180000);
+		if (client->mode != &idevicerestore_modes[MODE_RESTORE] || (client->flags & FLAG_QUIT)) {
+			mutex_unlock(&client->device_event_mutex);
+			error("ERROR: Device failed to enter restore mode.\n");
+			if (delete_fs && filesystem)
+				unlink(filesystem);
+			return -1;
+		}
 		mutex_unlock(&client->device_event_mutex);
-		error("ERROR: Device failed to enter restore mode.\n");
-		if (delete_fs && filesystem)
-			unlink(filesystem);
-		return -1;
 	}
-	mutex_unlock(&client->device_event_mutex);
 
 	// device is finally in restore mode, let's do this
 	if (client->mode->index == MODE_RESTORE) {
@@ -1418,6 +1432,9 @@ void idevicerestore_client_free(struct idevicerestore_client_t* client)
 	if (client->cache_dir) {
 		free(client->cache_dir);
 	}
+	if (client->root_ticket) {
+		free(client->root_ticket);
+	}
 	free(client);
 }
 
@@ -1528,7 +1545,7 @@ int main(int argc, char* argv[]) {
 		client->flags |= FLAG_INTERACTIVE;
 	}
 
-	while ((opt = getopt_long(argc, argv, "dhcesxtpli:u:nC:ky", longopts, &optindex)) > 0) {
+	while ((opt = getopt_long(argc, argv, "dhcesxtpli:u:nC:kyRT:", longopts, &optindex)) > 0) {
 		switch (opt) {
 		case 'h':
 			usage(argc, argv, 0);
@@ -1605,6 +1622,21 @@ int main(int argc, char* argv[]) {
 			client->flags &= ~FLAG_INTERACTIVE;
 			break;
 
+		case 'R':
+			client->flags |= FLAG_ALLOW_RESTORE_MODE;
+			break;
+
+		case 'T': {
+			size_t root_ticket_len = 0;
+			unsigned char* root_ticket = NULL;
+			if (read_file(optarg, (void**)&root_ticket, &root_ticket_len) != 0) {
+				return -1;
+			}
+			client->root_ticket = root_ticket;
+			client->root_ticket_len = (int)root_ticket_len;
+			info("Using ApTicket found at %s length %u\n", optarg, client->root_ticket_len);
+			break;
+		}
 		default:
 			usage(argc, argv, 1);
 			return -1;
@@ -1706,6 +1738,9 @@ int is_image4_supported(struct idevicerestore_client_t* client)
 	case MODE_NORMAL:
 		res = normal_is_image4_supported(client);
 		break;
+	case MODE_RESTORE:
+		res = restore_is_image4_supported(client);
+		break;
 	case MODE_DFU:
 		res = dfu_is_image4_supported(client);
 		break;
diff --git a/src/idevicerestore.h b/src/idevicerestore.h
index 16867b3..7d7fa53 100644
--- a/src/idevicerestore.h
+++ b/src/idevicerestore.h
@@ -43,6 +43,7 @@ extern "C" {
 #define FLAG_SHSHONLY        (1 << 7)
 #define FLAG_LATEST          (1 << 8)
 #define FLAG_INTERACTIVE     (1 << 9)
+#define FLAG_ALLOW_RESTORE_MODE (1 << 10)
 
 struct idevicerestore_client_t;
 
diff --git a/src/restore.c b/src/restore.c
index fea57fb..7a4e4fb 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -297,6 +297,46 @@ irecv_device_t restore_get_irecv_device(struct idevicerestore_client_t* client)
 	return irecv_device;
 }
 
+int restore_is_image4_supported(struct idevicerestore_client_t* client)
+{
+	int result = 0;
+	plist_t hwinfo = NULL;
+	idevice_t device = NULL;
+	restored_client_t restore = NULL;
+	restored_error_t restore_error = RESTORE_E_SUCCESS;
+
+	if (idevice_new(&device, client->udid) != IDEVICE_E_SUCCESS) {
+		error("ERROR: Could not connect to device %s\n", client->udid);
+		return -1;
+	}
+
+	restore_error = restored_client_new(device, &restore, "idevicerestore");
+	if (restore_error != RESTORE_E_SUCCESS) {
+		idevice_free(device);
+		return -1;
+	}
+
+	if (restored_query_type(restore, NULL, NULL) != RESTORE_E_SUCCESS) {
+		restored_client_free(restore);
+		idevice_free(device);
+		return -1;
+	}
+
+	restore_error = restored_query_value(restore, "HardwareInfo", &hwinfo);
+	if (restore_error == RESTORE_E_SUCCESS) {
+		uint8_t b = 0;
+		plist_t node = plist_dict_get_item(hwinfo, "SupportsImage4");
+		if (node && plist_get_node_type(node) == PLIST_BOOLEAN) {
+			plist_get_bool_val(node, &b);
+			result = b;
+		}
+	}
+	restored_client_free(restore);
+	idevice_free(device);
+
+	return result;
+}
+
 int restore_reboot(struct idevicerestore_client_t* client)
 {
 	if(client->restore == NULL) {
@@ -811,46 +851,51 @@ int restore_send_root_ticket(restored_client_t restore, struct idevicerestore_cl
 {
 	restored_error_t restore_error;
 	plist_t dict;
-	unsigned char* data = NULL;
-	unsigned int len = 0;
 
 	info("About to send RootTicket...\n");
 
-	if (!client->tss && !(client->flags & FLAG_CUSTOM)) {
-		error("ERROR: Cannot send RootTicket without TSS\n");
-		return -1;
-	}
+	if (client->root_ticket) {
+		dict = plist_new_dict();
+		plist_dict_set_item(dict, "RootTicketData", plist_new_data((char*)client->root_ticket, client->root_ticket_len));
+	} else {
+		unsigned char* data = NULL;
+		unsigned int len = 0;
 
-	if (client->image4supported) {
-		if (tss_response_get_ap_img4_ticket(client->tss, &data, &len) < 0) {
-			error("ERROR: Unable to get ApImg4Ticket from TSS\n");
+		if (!client->tss && !(client->flags & FLAG_CUSTOM)) {
+			error("ERROR: Cannot send RootTicket without TSS\n");
 			return -1;
 		}
-	} else {
-		if (!(client->flags & FLAG_CUSTOM) && (tss_response_get_ap_ticket(client->tss, &data, &len) < 0)) {
-			error("ERROR: Unable to get ticket from TSS\n");
-			return -1;
+
+		if (client->image4supported) {
+			if (tss_response_get_ap_img4_ticket(client->tss, &data, &len) < 0) {
+				error("ERROR: Unable to get ApImg4Ticket from TSS\n");
+				return -1;
+			}
+		} else {
+			if (!(client->flags & FLAG_CUSTOM) && (tss_response_get_ap_ticket(client->tss, &data, &len) < 0)) {
+				error("ERROR: Unable to get ticket from TSS\n");
+				return -1;
+			}
 		}
-	}
 
-	dict = plist_new_dict();
-	if (data && (len > 0)) {
-		plist_dict_set_item(dict, "RootTicketData", plist_new_data((char*)data, len));
-	} else {
-		info("NOTE: not sending RootTicketData (no data present)\n");
+		dict = plist_new_dict();
+		if (data && (len > 0)) {
+			plist_dict_set_item(dict, "RootTicketData", plist_new_data((char*)data, len));
+		} else {
+			info("NOTE: not sending RootTicketData (no data present)\n");
+		}
+		free(data);
 	}
 
 	info("Sending RootTicket now...\n");
 	restore_error = restored_send(restore, dict);
+	plist_free(dict);
 	if (restore_error != RESTORE_E_SUCCESS) {
 		error("ERROR: Unable to send RootTicket (%d)\n", restore_error);
-		plist_free(dict);
 		return -1;
 	}
 
 	info("Done sending RootTicket\n");
-	plist_free(dict);
-	free(data);
 	return 0;
 }
 
diff --git a/src/restore.h b/src/restore.h
index d4cfc46..bc41753 100644
--- a/src/restore.h
+++ b/src/restore.h
@@ -47,6 +47,7 @@ int restore_check_mode(struct idevicerestore_client_t* client);
 irecv_device_t restore_get_irecv_device(struct idevicerestore_client_t* client);
 int restore_client_new(struct idevicerestore_client_t* client);
 void restore_client_free(struct idevicerestore_client_t* client);
+int restore_is_image4_supported(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);