Implement handling of devices which do not require a tss request

This introduces passing around the "selected" build identity and
fixes code logic issues to make idevicerestore work again.

8 files changed, 184 insertions, 125 deletions

diff --git a/src/dfu.c b/src/dfu.c
index 6fd2648..0b3ad51 100644
--- a/src/dfu.c
+++ b/src/dfu.c
@@ -108,10 +108,11 @@ int dfu_check_mode() {
 	return -1;
 }
 
-int dfu_enter_recovery(struct idevicerestore_client_t* client) {
+int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_identity) {
 	irecv_client_t dfu = NULL;
 	const char* component = "iBSS";
 	irecv_error_t dfu_error = IRECV_E_SUCCESS;
+
 	if (recovery_open_with_timeout(client) < 0 || dfu->mode != kDfuMode) {
 		error("ERROR: Unable to connect to DFU device\n");
 		if (dfu)
@@ -119,7 +120,7 @@ int dfu_enter_recovery(struct idevicerestore_client_t* client) {
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, "iBSS") < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device\n", component);
 		irecv_close(dfu);
 		return -1;
diff --git a/src/dfu.h b/src/dfu.h
index 0e20448..259cb6a 100644
--- a/src/dfu.h
+++ b/src/dfu.h
@@ -37,7 +37,7 @@ struct dfu_client_t {
 
 int dfu_client_new(struct idevicerestore_client_t* client, uint32_t timeout);
 void dfu_client_free(struct idevicerestore_client_t* client);
-int dfu_enter_recovery(struct idevicerestore_client_t* client);
+int dfu_enter_recovery(struct idevicerestore_client_t* client, plist_t build_identity);
 
 
 #ifdef __cplusplus
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index 19ae8be..7982ed7 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -167,16 +167,15 @@ int main(int argc, char* argv[]) {
 	}
 
 	// choose whether this is an upgrade or a restore (default to upgrade)
-	plist_t tss = NULL;
+	client->tss = NULL;
 	plist_t build_identity = NULL;
 	if (client->flags & FLAG_ERASE) {
 		build_identity = get_build_identity(client, buildmanifest, 0);
 		if (build_identity == NULL) {
-			error("ERROR: Unable to find build any identities\n");
+			error("ERROR: Unable to find any build identities\n");
 			plist_free(buildmanifest);
 			return -1;
 		}
-
 	} else {
 		// loop through all build identities in the build manifest
 		// and list the valid ones
@@ -184,24 +183,33 @@ int main(int argc, char* argv[]) {
 		int valid_builds = 0;
 		int build_count = get_build_count(buildmanifest);
 		for (i = 0; i < build_count; i++) {
-			if (client->device->index > DEVICE_IPOD2G) {
-				build_identity = get_build_identity(client, buildmanifest, i);
-				if (get_shsh_blobs(client, ecid, build_identity, &tss) < 0) {
-					// if this fails then no SHSH blobs have been saved
-					// for this build identity, so check the next one
-					continue;
-				}
-				valid_builds++;
+			build_identity = get_build_identity(client, buildmanifest, i);
+			valid_builds++;
+		}
+	}
+
+	if (client->flags & FLAG_CUSTOM > 0) {
+		if (client->device->index > DEVICE_IPOD2G) {
+			if (get_shsh_blobs(client, ecid, build_identity, &client->tss) < 0) {
+				error("ERROR: Unable to get SHSH blobs for this device\n");
+				return -1;
 			}
 		}
+
+		/* verify if we have tss records if required */
+		if ((client->device->index > DEVICE_IPOD2G) && (client->tss == NULL)) {
+			error("ERROR: Unable to proceed without a tss record.\n");
+			plist_free(buildmanifest);
+			return -1;
+		}
 	}
 
 	// Extract filesystem from IPSW and return its name
 	char* filesystem = NULL;
-	if (extract_filesystem(client, ipsw, build_identity, &filesystem) < 0) {
+	if (extract_filesystem(client, client->ipsw, build_identity, &filesystem) < 0) {
 		error("ERROR: Unable to extract filesystem from IPSW\n");
-		if (tss)
-			plist_free(tss);
+		if (client->tss)
+			plist_free(client->tss);
 		plist_free(buildmanifest);
 		return -1;
 	}
@@ -211,8 +219,8 @@ int main(int argc, char* argv[]) {
 		info("Entering recovery mode...\n");
 		if (normal_enter_recovery(client) < 0) {
 			error("ERROR: Unable to place device into recovery mode\n");
-			if (tss)
-				plist_free(tss);
+			if (client->tss)
+				plist_free(client->tss);
 			plist_free(buildmanifest);
 			return -1;
 		}
@@ -220,22 +228,22 @@ int main(int argc, char* argv[]) {
 
 	// if the device is in DFU mode, place device into recovery mode
 	if (client->mode->index == MODE_DFU) {
-		if (dfu_enter_recovery(client) < 0) {
+		if (dfu_enter_recovery(client, build_identity) < 0) {
 			error("ERROR: Unable to place device into recovery mode\n");
 			plist_free(buildmanifest);
-			if (tss)
-				plist_free(tss);
+			if (client->tss)
+				plist_free(client->tss);
 			return -1;
 		}
 	}
 
 	// if the device is in recovery mode, place device into restore mode
 	if (client->mode->index == MODE_RECOVERY) {
-		if (recovery_enter_restore(uuid, ipsw, tss) < 0) {
+		if (recovery_enter_restore(client, build_identity) < 0) {
 			error("ERROR: Unable to place device into restore mode\n");
 			plist_free(buildmanifest);
-			if (tss)
-				plist_free(tss);
+			if (client->tss)
+				plist_free(client->tss);
 			return -1;
 		}
 	}
@@ -243,7 +251,7 @@ int main(int argc, char* argv[]) {
 	// device is finally in restore mode, let's do this
 	if (client->mode->index == MODE_RESTORE) {
 		info("Restoring device... \n");
-		if (restore_device(client, uuid, ipsw, tss, filesystem) < 0) {
+		if (restore_device(client, build_identity, filesystem) < 0) {
 			error("ERROR: Unable to restore device\n");
 			return -1;
 		}
@@ -545,7 +553,7 @@ int extract_filesystem(struct idevicerestore_client_t* client, const char* ipsw,
 	return 0;
 }
 
-int get_signed_component(struct idevicerestore_client_t* client, const char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size) {
+int ipsw_get_component_by_path(const char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size) {
 	img3_file* img3 = NULL;
 	uint32_t component_size = 0;
 	char* component_data = NULL;
@@ -564,36 +572,40 @@ int get_signed_component(struct idevicerestore_client_t* client, const char* ips
 		return -1;
 	}
 
-	img3 = img3_parse_file(component_data, component_size);
-	if (img3 == NULL) {
-		error("ERROR: Unable to parse IMG3: %s\n", component_name);
+	if (tss) {
+		info("Signing img3...\n");
+		img3 = img3_parse_file(component_data, component_size);
+		if (img3 == NULL) {
+			error("ERROR: Unable to parse IMG3: %s\n", component_name);
+			free(component_data);
+			return -1;
+		}
 		free(component_data);
-		return -1;
-	}
-	free(component_data);
 
-	if (tss_get_blob_by_path(tss, path, &component_blob) < 0) {
-		error("ERROR: Unable to get SHSH blob for TSS %s entry\n", component_name);
-		img3_free(img3);
-		return -1;
-	}
+		/* sign the blob if required */
+		if (tss_get_blob_by_path(tss, path, &component_blob) < 0) {
+			error("ERROR: Unable to get SHSH blob for TSS %s entry\n", component_name);
+			img3_free(img3);
+			return -1;
+		}
 
-	if (client->device->index > DEVICE_IPOD2G && (client->flags & FLAG_CUSTOM) == 0) {
 		if (img3_replace_signature(img3, component_blob) < 0) {
 			error("ERROR: Unable to replace IMG3 signature\n");
 			free(component_blob);
 			img3_free(img3);
 			return -1;
 		}
-	}
-	free(component_blob);
 
-	if (img3_get_data(img3, &component_data, &component_size) < 0) {
-		error("ERROR: Unable to reconstruct IMG3\n");
+		if (component_blob)
+			free(component_blob);
+
+		if (img3_get_data(img3, &component_data, &component_size) < 0) {
+			error("ERROR: Unable to reconstruct IMG3\n");
+			img3_free(img3);
+			return -1;
+		}
 		img3_free(img3);
-		return -1;
 	}
-	img3_free(img3);
 
 	if (idevicerestore_debug) {
 		write_file(component_name, component_data, component_size);
diff --git a/src/idevicerestore.h b/src/idevicerestore.h
index 3213b0c..f529b5b 100644
--- a/src/idevicerestore.h
+++ b/src/idevicerestore.h
@@ -42,7 +42,7 @@ int extract_buildmanifest(struct idevicerestore_client_t* client, const char* ip
 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 get_signed_component(struct idevicerestore_client_t* client, const char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size);
+int ipsw_get_component_by_path(const char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size);
 int build_identity_get_component_path(plist_t build_identity, const char* component, char** path);
 
 #ifdef __cplusplus
diff --git a/src/recovery.c b/src/recovery.c
index 233224a..6a38343 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -45,6 +45,7 @@ int recovery_client_new(struct idevicerestore_client_t* client) {
 		error("ERROR: Out of memory\n");
 		return -1;
 	}
+
 	client->recovery = recovery;
 
 	if (recovery_open_with_timeout(client) < 0) {
@@ -122,28 +123,28 @@ int recovery_check_mode() {
 	return 0;
 }
 
-int recovery_enter_restore(struct idevicerestore_client_t* client) {
+int recovery_enter_restore(struct idevicerestore_client_t* client, plist_t build_identity) {
 	idevice_t device = NULL;
 	restored_client_t restore = NULL;
 
 	// upload data to make device boot restore mode
-	if (recovery_send_ibec(client) < 0) {
+	if (recovery_send_ibec(client, build_identity) < 0) {
 		error("ERROR: Unable to send iBEC\n");
 		return -1;
 	}
-	sleep(1);
+	sleep(2);
 
-	if (recovery_send_applelogo(client) < 0) {
+	if (recovery_send_applelogo(client, build_identity) < 0) {
 		error("ERROR: Unable to send AppleLogo\n");
 		return -1;
 	}
 
-	if (recovery_send_devicetree(client) < 0) {
+	if (recovery_send_devicetree(client, build_identity) < 0) {
 		error("ERROR: Unable to send DeviceTree\n");
 		return -1;
 	}
 
-	if (recovery_send_ramdisk(client) < 0) {
+	if (recovery_send_ramdisk(client, build_identity) < 0) {
 		error("ERROR: Unable to send Ramdisk\n");
 		return -1;
 	}
@@ -154,7 +155,10 @@ int recovery_enter_restore(struct idevicerestore_client_t* client) {
 	printf("Hit any key to continue...");
 	getchar();
 
-	if (recovery_send_kernelcache(client) < 0) {
+	info("Resetting recovery mode connection...\n");
+	irecv_reset(client->recovery->client);
+
+	if (recovery_send_kernelcache(client, build_identity) < 0) {
 		error("ERROR: Unable to send KernelCache\n");
 		return -1;
 	}
@@ -170,29 +174,42 @@ int recovery_enter_restore(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-int recovery_send_signed_component(struct idevicerestore_client_t* client, const char* component) {
-	int size = 0;
+int recovery_send_component(struct idevicerestore_client_t* client, plist_t build_identity, const char* component) {
+	uint32_t size = 0;
 	char* data = NULL;
 	char* path = NULL;
 	char* blob = NULL;
 	irecv_error_t error = 0;
 
-	if (tss_get_entry_path(client->tss, component, &path) < 0) {
-		error("ERROR: Unable to get component path\n");
-		return -1;
+	if (client->tss) {
+		if (tss_get_entry_path(client->tss, component, &path) < 0) {
+			error("ERROR: Unable to get component path\n");
+			return -1;
+		}
+	} else {
+		if (build_identity_get_component_path(build_identity, component, &path) < 0) {
+			error("ERROR: Unable to get component: %s\n", component);
+			if (path)
+				free(path);
+			return -1;
+		}
 	}
 
-	if (get_signed_component(client, client->ipsw, client->tss, path, &data, &size) < 0) {
-		error("ERROR: Unable to get signed component: %s\n", component);
+	info("Resetting recovery mode connection...\n");
+	irecv_reset(client->recovery->client);
+
+	if (ipsw_get_component_by_path(client->ipsw, client->tss, path, &data, &size) < 0) {
+		error("ERROR: Unable to get component: %s\n", component);
 		free(path);
 		return -1;
 	}
-	free(path);
 
-	info("Sending %s...\n", component);
+	info("Sending %s (%d bytes)...\n", component, size);
+
 	error = irecv_send_buffer(client->recovery->client, data, size);
+	free(path);
 	if (error != IRECV_E_SUCCESS) {
-		error("ERROR: Unable to send component: %s\n", component);
+		error("ERROR: Unable to send %s component: %s\n", component, irecv_strerror(error));
 		free(data);
 		return -1;
 	}
@@ -201,10 +218,8 @@ int recovery_send_signed_component(struct idevicerestore_client_t* client, const
 	return 0;
 }
 
-int recovery_send_ibec(struct idevicerestore_client_t* client) {
-	const char* component = "iBEC";
+static int recovery_enable_autoboot(struct idevicerestore_client_t* client) {
 	irecv_error_t recovery_error = IRECV_E_SUCCESS;
-
 	//recovery_error = irecv_send_command(client->recovery->client, "setenv auto-boot true");
 	recovery_error = irecv_setenv(client->recovery->client, "auto-boot", "true");
 	if (recovery_error != IRECV_E_SUCCESS) {
@@ -218,7 +233,18 @@ int recovery_send_ibec(struct idevicerestore_client_t* client) {
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, "iBEC") < 0) {
+	return 0;
+}
+
+int recovery_send_ibec(struct idevicerestore_client_t* client, plist_t build_identity) {
+	const char* component = "iBEC";
+	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+
+	if (recovery_enable_autoboot(client) < 0) {
+		return -1;
+	}
+
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
 		return -1;
 	}
@@ -232,9 +258,8 @@ int recovery_send_ibec(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-int recovery_send_applelogo(struct idevicerestore_client_t* client) {
-	irecv_client_t recovery = NULL;
-	const char* component = "applelogo";
+int recovery_send_applelogo(struct idevicerestore_client_t* client, plist_t build_identity) {
+	const char* component = "AppleLogo";
 	irecv_error_t recovery_error = IRECV_E_SUCCESS;
 
 	info("Sending %s...\n", component);
@@ -242,40 +267,37 @@ int recovery_send_applelogo(struct idevicerestore_client_t* client) {
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, "AppleLogo") < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
-		irecv_close(recovery);
 		return -1;
 	}
 
-	recovery_error = irecv_send_command(recovery, "setpicture 1");
+	recovery_error = irecv_send_command(client->recovery->client, "setpicture 1");
 	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to set %s\n", component);
-		irecv_close(recovery);
 		return -1;
 	}
 
-	recovery_error = irecv_send_command(recovery, "bgcolor 0 0 0");
+	recovery_error = irecv_send_command(client->recovery->client, "bgcolor 0 0 0");
 	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to display %s\n", component);
-		irecv_close(recovery);
 		return -1;
 	}
 
-	irecv_close(recovery);
-	recovery = NULL;
 	return 0;
 }
 
-int recovery_send_devicetree(struct idevicerestore_client_t* client) {
-	const char* component = "devicetree";
+int recovery_send_devicetree(struct idevicerestore_client_t* client, plist_t build_identity) {
+	const char* component = "RestoreDeviceTree";
 	irecv_error_t recovery_error = IRECV_E_SUCCESS;
 
-	if (recovery_open_with_timeout(client) < 0) {
-		return -1;
+	if(client->recovery == NULL) {
+		if (recovery_open_with_timeout(client) < 0) {
+			return -1;
+		}
 	}
 
-	if (recovery_send_signed_component(client, "RestoreDeviceTree") < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
 		return -1;
 	}
@@ -289,16 +311,17 @@ int recovery_send_devicetree(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-int recovery_send_ramdisk(struct idevicerestore_client_t* client) {
+int recovery_send_ramdisk(struct idevicerestore_client_t* client, plist_t build_identity) {
+	const char *component = "RestoreRamDisk";
 	irecv_error_t recovery_error = IRECV_E_SUCCESS;
-	const char *component = "ramdisk";
 
-	recovery_error = recovery_open_with_timeout(client);
-	if (recovery_error != IRECV_E_SUCCESS) {
-		return -1;
+	if(client->recovery == NULL) {
+		if (recovery_open_with_timeout(client) < 0) {
+			return -1;
+		}
 	}
 
-	if (recovery_send_signed_component(client, "RestoreRamDisk") < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
 		return -1;
 	}
@@ -312,15 +335,15 @@ int recovery_send_ramdisk(struct idevicerestore_client_t* client) {
 	return 0;
 }
 
-int recovery_send_kernelcache(struct idevicerestore_client_t* client) {
-	const char* component = "kernelcache";
+int recovery_send_kernelcache(struct idevicerestore_client_t* client, plist_t build_identity) {
+	const char* component = "RestoreKernelCache";
 	irecv_error_t recovery_error = IRECV_E_SUCCESS;
 
 	if (recovery_open_with_timeout(client) < 0) {
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, "RestoreKernelCache") < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
 		return -1;
 	}
@@ -337,8 +360,10 @@ int recovery_send_kernelcache(struct idevicerestore_client_t* client) {
 int recovery_get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
 	irecv_error_t recovery_error = IRECV_E_SUCCESS;
 
-	if (recovery_open_with_timeout(client) < 0) {
-		return -1;
+	if(client->recovery == NULL) {
+		if (recovery_open_with_timeout(client) < 0) {
+			return -1;
+		}
 	}
 
 	recovery_error = irecv_get_ecid(client->recovery->client, ecid);
diff --git a/src/recovery.h b/src/recovery.h
index 5d1129f..b7cc0e4 100644
--- a/src/recovery.h
+++ b/src/recovery.h
@@ -43,12 +43,12 @@ int recovery_check_mode();
 int recovery_open_with_timeout(struct idevicerestore_client_t* client);
 int recovery_client_new(struct idevicerestore_client_t* client);
 void recovery_client_free(struct idevicerestore_client_t* client);
-int recovery_send_signed_component(struct idevicerestore_client_t* client, const char* component);
-int recovery_send_ibec(struct idevicerestore_client_t* client);
-int recovery_send_applelogo(struct idevicerestore_client_t* client);
-int recovery_send_devicetree(struct idevicerestore_client_t* client);
-int recovery_send_ramdisk(struct idevicerestore_client_t* client);
-int recovery_send_kernelcache(struct idevicerestore_client_t* client);
+int recovery_send_component(struct idevicerestore_client_t* client, plist_t build_identity, const char* component);
+int recovery_send_ibec(struct idevicerestore_client_t* client, plist_t build_identity);
+int recovery_send_applelogo(struct idevicerestore_client_t* client, plist_t build_identity);
+int recovery_send_devicetree(struct idevicerestore_client_t* client, plist_t build_identity);
+int recovery_send_ramdisk(struct idevicerestore_client_t* client, plist_t build_identity);
+int recovery_send_kernelcache(struct idevicerestore_client_t* client, plist_t build_identity);
 int recovery_get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid);
 int recovery_get_cpid(struct idevicerestore_client_t* client, uint32_t* cpid);
 int recovery_get_bdid(struct idevicerestore_client_t* client, uint32_t* bdid);
diff --git a/src/restore.c b/src/restore.c
index bf4b62f..5d45296 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -186,7 +186,8 @@ int restore_reboot(struct idevicerestore_client_t* client) {
 	idevice_t device = NULL;
 	restored_client_t restore = NULL;
 	restored_error_t restore_error = RESTORE_E_SUCCESS;
-	if(!client->restore) {
+
+	if(client->restore == NULL) {
 		if (restore_open_with_timeout(client) < 0) {
 			error("ERROR: Unable to open device in restore mode\n");
 			return -1;
@@ -398,7 +399,7 @@ int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	return 0;
 }
 
-int restore_send_kernelcache(restored_client_t client, const char* ipsw, plist_t tss) {
+int restore_send_kernelcache(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity) {
 	int size = 0;
 	char* data = NULL;
 	char* path = NULL;
@@ -407,22 +408,31 @@ int restore_send_kernelcache(restored_client_t client, const char* ipsw, plist_t
 	restored_error_t restore_error = RESTORE_E_SUCCESS;
 
 	info("Sending kernelcache\n");
-	if (tss_get_entry_path(tss, "KernelCache", &path) < 0) {
-		error("ERROR: Unable to find kernelcache path\n");
-		return -1;
+
+	if (client->tss) {
+		if (tss_get_entry_path(client->tss, "KernelCache", &path) < 0) {
+			error("ERROR: Unable to get KernelCache path\n");
+			return -1;
+		}
+	} else {
+		if (build_identity_get_component_path(build_identity, "KernelCache", &path) < 0) {
+			error("ERROR: Unable to find kernelcache path\n");
+			if (path)
+				free(path);
+			return -1;
+		}
 	}
 
-	if (get_signed_component(client, ipsw, tss, path, &data, &size) < 0) {
+	if (ipsw_get_component_by_path(client->ipsw, client->tss, path, &data, &size) < 0) {
 		error("ERROR: Unable to get kernelcache file\n");
 		return -1;
 	}
 
-
 	dict = plist_new_dict();
 	blob = plist_new_data(data, size);
 	plist_dict_insert_item(dict, "KernelCacheFile", blob);
 
-	restore_error = restored_send(client, dict);
+	restore_error = restored_send(restore, dict);
 	if (restore_error != RESTORE_E_SUCCESS) {
 		error("ERROR: Unable to send kernelcache data\n");
 		plist_free(dict);
@@ -435,7 +445,7 @@ int restore_send_kernelcache(restored_client_t client, const char* ipsw, plist_t
 	return 0;
 }
 
-int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss) {
+int restore_send_nor(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity) {
 	char* llb_path = NULL;
 	char* llb_filename = NULL;
 	char firmware_path[256];
@@ -452,9 +462,18 @@ int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss) {
 	plist_t norimage_array = NULL;
 	restored_error_t ret = RESTORE_E_SUCCESS;
 
-	if (tss_get_entry_path(tss, "LLB", &llb_path) < 0) {
-		error("ERROR: Unable to get LLB info from TSS response\n");
-		return -1;
+	if (client->tss) {
+		if (tss_get_entry_path(client->tss, "LLB", &llb_path) < 0) {
+			error("ERROR: Unable to get LLB path\n");
+			return -1;
+		}
+	} else {
+		if (build_identity_get_component_path(build_identity, "LLB", &llb_path) < 0) {
+			error("ERROR: Unable to get component: LLB\n");
+			if (llb_path)
+				free(llb_path);
+			return -1;
+		}
 	}
 
 	llb_filename = strstr(llb_path, "LLB");
@@ -472,7 +491,7 @@ int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss) {
 	snprintf(manifest_file, sizeof(manifest_file), "%s/manifest", firmware_path);
 	info("Getting firmware manifest %s\n", manifest_file);
 
-	if (ipsw_extract_to_memory(ipsw, manifest_file, &manifest_data, &manifest_size) < 0) {
+	if (ipsw_extract_to_memory(client->ipsw, manifest_file, &manifest_data, &manifest_size) < 0) {
 		error("ERROR: Unable to extract firmware manifest from ipsw\n");
 		free(llb_path);
 		return -1;
@@ -485,7 +504,7 @@ int restore_send_nor(restored_client_t client, const 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(client, ipsw, tss, firmware_filename, &llb_data, &llb_size) < 0) {
+		if (ipsw_get_component_by_path(client->ipsw, client->tss, firmware_filename, &llb_data, &llb_size) < 0) {
 			error("ERROR: Unable to get signed LLB\n");
 			return -1;
 		}
@@ -498,7 +517,7 @@ int restore_send_nor(restored_client_t client, const 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(client, ipsw, tss, firmware_filename, &nor_data, &nor_size) < 0) {
+		if (ipsw_get_component_by_path(client->ipsw, client->tss, firmware_filename, &nor_data, &nor_size) < 0) {
 			error("ERROR: Unable to get signed firmware %s\n", firmware_filename);
 			break;
 		}
@@ -513,7 +532,7 @@ int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss) {
 
 	debug_plist(dict);
 
-	ret = restored_send(client, dict);
+	ret = restored_send(restore, dict);
 	if (ret != RESTORE_E_SUCCESS) {
 		error("ERROR: Unable to send kernelcache data\n");
 		plist_free(dict);
@@ -524,7 +543,7 @@ int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss) {
 	return 0;
 }
 
-int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idevice_t device, restored_client_t restore, plist_t message, plist_t tss, const char* ipsw, const char* filesystem) {
+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) {
 	char* type = NULL;
 	plist_t node = NULL;
 
@@ -543,7 +562,7 @@ int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idev
 		}
 
 		else if (!strcmp(type, "KernelCache")) {
-			if(restore_send_kernelcache(restore, ipsw, tss) < 0) {
+			if(restore_send_kernelcache(restore, client, build_identity) < 0) {
 				error("ERROR: Unable to send kernelcache\n");
 				return -1;
 			}
@@ -551,7 +570,7 @@ int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idev
 
 		else if (!strcmp(type, "NORData")) {
 			if(client->flags & FLAG_EXCLUDE > 0) {
-				if(restore_send_nor(restore, ipsw, tss) < 0) {
+				if(restore_send_nor(restore, client, build_identity) < 0) {
 					error("ERROR: Unable to send NOR data\n");
 					return -1;
 				}
@@ -568,7 +587,7 @@ int restore_handle_data_request_msg(struct idevicerestore_client_t* client, idev
 	return 0;
 }
 
-int restore_device(struct idevicerestore_client_t* client, const char* uuid, const char* ipsw, plist_t tss, const char* filesystem) {
+int restore_device(struct idevicerestore_client_t* client, plist_t build_identity, const char* filesystem) {
 	int error = 0;
 	char* type = NULL;
 	char* kernel = NULL;
@@ -586,6 +605,8 @@ int restore_device(struct idevicerestore_client_t* client, const char* uuid, con
 	}
 	info("Device has successfully entered restore mode\n");
 
+	restore = client->restore->client;
+
 	// start the restore process
 	restore_error = restored_start_restore(restore);
 	if (restore_error != RESTORE_E_SUCCESS) {
@@ -619,7 +640,7 @@ int restore_device(struct idevicerestore_client_t* client, const char* uuid, con
 		// files sent to the server by the client. these data requests include
 		// SystemImageData, KernelCache, and NORData requests
 		if (!strcmp(type, "DataRequestMsg")) {
-			error = restore_handle_data_request_msg(client, device, restore, message, tss, ipsw, filesystem);
+			error = restore_handle_data_request_msg(client, device, restore, message, build_identity, filesystem);
 		}
 
 		// progress notification messages sent by the restored inform the client
diff --git a/src/restore.h b/src/restore.h
index 5446aa8..9c11c34 100644
--- a/src/restore.h
+++ b/src/restore.h
@@ -47,10 +47,10 @@ 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_data_request_msg(struct idevicerestore_client_t* client, idevice_t device, restored_client_t restore, plist_t message, plist_t tss, const char* ipsw, const char* filesystem);
-int restore_send_nor(restored_client_t client, const char* ipsw, plist_t tss);
-int restore_send_kernelcache(restored_client_t client, const char* ipsw, plist_t tss);
-int restore_device(struct idevicerestore_client_t* client, const char* uuid, const char* ipsw, plist_t tss, const char* filesystem);
+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_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);