Use DeviceGeneratedRequest plist for SE TSS requests

2 files changed, 22 insertions, 8 deletions

diff --git a/src/restore.c b/src/restore.c
index 69f2f09..486ffb5 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -2105,7 +2105,7 @@ static int restore_send_image_data(restored_client_t restore, struct idevicerest
 	return 0;
 }
 
-static plist_t restore_get_se_firmware_data(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity, plist_t p_info)
+static plist_t restore_get_se_firmware_data(restored_client_t restore, struct idevicerestore_client_t* client, plist_t build_identity, plist_t p_info, plist_t arguments)
 {
 	const char *comp_name = NULL;
 	char *comp_path = NULL;
@@ -2114,6 +2114,7 @@ static plist_t restore_get_se_firmware_data(restored_client_t restore, struct id
 	plist_t parameters = NULL;
 	plist_t request = NULL;
 	plist_t response = NULL;
+	plist_t p_dgr = NULL;
 	int ret;
 	uint64_t chip_id = 0;
 	plist_t node = plist_dict_get_item(p_info, "SE,ChipID");
@@ -2137,6 +2138,14 @@ static plist_t restore_get_se_firmware_data(restored_client_t restore, struct id
 		debug("DEBUG: %s: using %s\n", __func__, comp_name);
 	}
 
+	p_dgr = plist_dict_get_item(arguments, "DeviceGeneratedRequest");
+	if (!p_dgr) {
+		info("NOTE: %s: No DeviceGeneratedRequest in firmware updater data request. Continuing anyway.\n", __func__);
+	} else if (!PLIST_IS_DICT(p_dgr)) {
+		error("ERROR: %s: DeviceGeneratedRequest has invalid type!\n", __func__);
+		return NULL;
+	}
+
 	if (build_identity_get_component_path(build_identity, comp_name, &comp_path) < 0) {
 		error("ERROR: Unable to get path for '%s' component\n", comp_name);
 		return NULL;
@@ -2167,7 +2176,7 @@ static plist_t restore_get_se_firmware_data(restored_client_t restore, struct id
 	plist_dict_merge(&parameters, p_info);
 
 	/* add required tags for SE TSS request */
-	tss_request_add_se_tags(request, parameters, NULL);
+	tss_request_add_se_tags(request, parameters, p_dgr);
 
 	plist_free(parameters);
 
@@ -2180,10 +2189,12 @@ static plist_t restore_get_se_firmware_data(restored_client_t restore, struct id
 		return NULL;
 	}
 
-	if (plist_dict_get_item(response, "SE,Ticket")) {
-		info("Received SE ticket\n");
+	if (plist_dict_get_item(response, "SE2,Ticket")) {
+		info("Received SE2,Ticket\n");
+	} else if (plist_dict_get_item(response, "SE,Ticket")) {
+		info("Received SE,Ticket\n");
 	} else {
-		error("ERROR: No 'SE,Ticket' in TSS response, this might not work\n");
+		error("ERROR: No 'SE ticket' in TSS response, this might not work\n");
 	}
 
 	plist_dict_set_item(response, "FirmwareData", plist_new_data((char*)component_data, (uint64_t) component_size));
@@ -3025,7 +3036,7 @@ static int restore_send_firmware_updater_data(restored_client_t restore, struct
 	plist_get_string_val(p_updater_name, &s_updater_name);
 
 	if (strcmp(s_updater_name, "SE") == 0) {
-		fwdict = restore_get_se_firmware_data(restore, client, build_identity, p_info);
+		fwdict = restore_get_se_firmware_data(restore, client, build_identity, p_info, arguments);
 		if (fwdict == NULL) {
 			error("ERROR: %s: Couldn't get SE firmware data\n", __func__);
 			goto error_out;
diff --git a/src/tss.c b/src/tss.c
index e1cf1f4..9fb74af 100644
--- a/src/tss.c
+++ b/src/tss.c
@@ -786,9 +786,7 @@ int tss_request_add_se_tags(plist_t request, plist_t parameters, plist_t overrid
 		return -1;
 	}
 
-	/* add tags indicating we want to get the SE,Ticket */
 	plist_dict_set_item(request, "@BBTicket", plist_new_bool(1));
-	plist_dict_set_item(request, "@SE,Ticket", plist_new_bool(1));
 
 	if (_plist_dict_copy_uint(request, parameters, "SE,ChipID", NULL) < 0) {
 		error("ERROR: %s: Unable to find required SE,ChipID in parameters\n", __func__);
@@ -864,6 +862,11 @@ int tss_request_add_se_tags(plist_t request, plist_t parameters, plist_t overrid
 		plist_dict_merge(&request, overrides);
 	}
 
+	/* fallback in case no @SE2,Ticket or @SE,Ticket was provided */
+	if (!plist_dict_get_item(request, "@SE2,Ticket") && !plist_dict_get_item(request, "@SE,Ticket")) {
+		plist_dict_set_item(request, "@SE,Ticket", plist_new_bool(1));
+	}
+
 	return 0;
 }