img4: Add support for stitching with additional TBM data

3 files changed, 191 insertions, 7 deletions

diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index 308c6cb..62fbc09 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -2566,9 +2566,9 @@ int personalize_component(const char *component_name, const unsigned char* compo
 	unsigned char* stitched_component = NULL;
 	unsigned int stitched_component_size = 0;
 
-	if (tss_response && tss_response_get_ap_img4_ticket(tss_response, &component_blob, &component_blob_size) == 0) {
+	if (tss_response && plist_dict_get_item(tss_response, "ApImg4Ticket")) {
 		/* stitch ApImg4Ticket into IMG4 file */
-		img4_stitch_component(component_name, component_data, component_size, component_blob, component_blob_size, &stitched_component, &stitched_component_size);
+		img4_stitch_component(component_name, component_data, component_size, tss_response, &stitched_component, &stitched_component_size);
 	} else {
 		/* try to get blob for current component from tss response */
 		if (tss_response && tss_response_get_blob_by_entry(tss_response, component_name, &component_blob) < 0) {
diff --git a/src/img4.c b/src/img4.c
index 6cff349..3a74c5e 100644
--- a/src/img4.c
+++ b/src/img4.c
@@ -24,6 +24,7 @@
 
 #include "common.h"
 #include "img4.h"
+#include "tss.h"
 
 #define ASN1_PRIVATE 0xc0
 #define ASN1_PRIMITIVE_TAG 0x1f
@@ -201,7 +202,7 @@ static void asn1_write_element(unsigned char **p, unsigned int *length, unsigned
 		}
 	}	break;
 	default:
-		fprintf(stderr, "ERROR: %s: type %02x is not implemented", __func__, type);
+		fprintf(stderr, "ERROR: %s: type %02x is not implemented\n", __func__, type);
 		return;
 	}
 }
@@ -393,7 +394,38 @@ static const char *_img4_get_component_tag(const char *compname)
 	return NULL;
 }
 
-int img4_stitch_component(const char* component_name, const unsigned char* component_data, unsigned int component_size, const unsigned char* blob, unsigned int blob_size, unsigned char** img4_data, unsigned int *img4_size)
+static void hexdump(char* data, int length)
+{
+        int i;
+        int j;
+        unsigned char c;
+
+                for (i = 0; i < length; i += 16) {
+                        fprintf(stderr, "%04x: ", i);
+                        for (j = 0; j < 16; j++) {
+                                if (i + j >= length) {
+                                        fprintf(stderr, "   ");
+                                        continue;
+                                }
+                                fprintf(stderr, "%02x ", *(data + i + j) & 0xff);
+                        }
+                        fprintf(stderr, "  | ");
+                        for (j = 0; j < 16; j++) {
+                                if (i + j >= length)
+                                        break;
+                                c = *(data + i + j);
+                                if ((c < 32) || (c > 127)) {
+                                        fprintf(stderr, ".");
+                                        continue;
+                                }
+                                fprintf(stderr, "%c", c);
+                        }
+                        fprintf(stderr, "\n");
+                }
+                fprintf(stderr, "\n");
+}
+
+int img4_stitch_component(const char* component_name, const unsigned char* component_data, unsigned int component_size, plist_t tss_response, unsigned char** img4_data, unsigned int *img4_size)
 {
 	unsigned char* magic_header = NULL;
 	unsigned int magic_header_size = 0;
@@ -404,8 +436,15 @@ int img4_stitch_component(const char* component_name, const unsigned char* compo
 	unsigned int content_size;
 	unsigned char* outbuf;
 	unsigned char* p;
+	unsigned char* blob = NULL;
+	unsigned int blob_size = 0;
 
-	if (!component_name || !component_data || component_size == 0 || !blob || blob_size == 0 || !img4_data || !img4_size) {
+	if (!component_name || !component_data || component_size == 0 || !tss_response || !img4_data || !img4_size) {
+		return -1;
+	}
+
+	if (tss_response_get_ap_img4_ticket(tss_response, &blob, &blob_size) != 0) {
+		error("ERROR: %s: Failed to get ApImg4Ticket from TSS response\n", __func__);
 		return -1;
 	}
 
@@ -435,13 +474,152 @@ int img4_stitch_component(const char* component_name, const unsigned char* compo
 		}
 	}
 
+	// check if we have a *-TBM entry for the given component
+	unsigned char *additional_data = NULL;
+	unsigned int additional_size = 0;
+	char *tbm_key = malloc(strlen(component_name) + 5);
+	sprintf(tbm_key, "%s-TBM", component_name);
+	plist_t tbm_dict = plist_dict_get_item(tss_response, tbm_key);
+	free(tbm_key);
+	if (tbm_dict) {
+		plist_t dt = plist_dict_get_item(tbm_dict, "ucon");
+		if (!dt) {
+			error("ERROR: %s: Missing ucon node in %s-TBM dictionary\n", __func__, component_name);
+			return -1;
+		}
+		uint64_t ucon_size = 0;
+		const char* ucon_data = plist_get_data_ptr(dt, &ucon_size);
+		if (!ucon_data) {
+			error("ERROR: %s: Missing ucon data in %s-TBM dictionary\n", __func__, component_name);
+			return -1;
+		}
+		dt = plist_dict_get_item(tbm_dict, "ucer");
+		if (!dt) {
+			error("ERROR: %s: Missing ucer data node in %s-TBM dictionary\n", __func__, component_name);
+			return -1;
+		}
+		uint64_t ucer_size = 0;
+		const char* ucer_data = plist_get_data_ptr(dt, &ucer_size);
+		if (!ucer_data) {
+			error("ERROR: %s: Missing ucer data in %s-TBM dictionary\n", __func__, component_name);
+			return -1;
+		}
+
+		unsigned char *im4rset = (unsigned char*)malloc(16 + 8 + 8 + ucon_size + 16 + 8 + 8 + ucer_size + 16);
+		unsigned char *p_im4rset = im4rset;
+		unsigned int im4rlen = 0;
+
+		// ----------- ucon ------------
+		// write priv ucon element
+		asn1_write_priv_element(&p_im4rset, &im4rlen, *(uint32_t*)"nocu");
+
+		// write ucon IA5STRING and ucon data
+		unsigned char ucon_seq[16];
+		unsigned char *p_ucon_seq = &ucon_seq[0];
+		unsigned int ucon_seq_hdr_len = 0;
+		asn1_write_element(&p_ucon_seq, &ucon_seq_hdr_len, ASN1_IA5_STRING, (void*)"ucon", -1);
+		asn1_write_element_header(ASN1_OCTET_STRING, ucon_size, &p_ucon_seq, &ucon_seq_hdr_len);
+
+		// write ucon sequence
+		unsigned char elem_seq[8];
+		unsigned char *p = &elem_seq[0];
+		unsigned int seq_hdr_len = 0;
+		asn1_write_element_header(ASN1_SEQUENCE | ASN1_CONSTRUCTED, ucon_seq_hdr_len + ucon_size, &p, &seq_hdr_len);
+
+		// add size to priv ucon element
+		asn1_write_size(ucon_seq_hdr_len + ucon_size + seq_hdr_len, &p_im4rset, &im4rlen);
+
+		// put it together
+		memcpy(p_im4rset, elem_seq, seq_hdr_len);
+		p_im4rset += seq_hdr_len;
+		im4rlen += seq_hdr_len;
+		memcpy(p_im4rset, ucon_seq, ucon_seq_hdr_len);
+		p_im4rset += ucon_seq_hdr_len;
+		im4rlen += ucon_seq_hdr_len;
+		memcpy(p_im4rset, ucon_data, ucon_size);
+		p_im4rset += ucon_size;
+		im4rlen += ucon_size;
+
+		// ----------- ucer ------------
+		// write priv ucer element
+		asn1_write_priv_element(&p_im4rset, &im4rlen, *(uint32_t*)"recu");
+
+		// write ucon IA5STRING and ucer data
+		unsigned char ucer_seq[16];
+		unsigned char *p_ucer_seq = &ucer_seq[0];
+		unsigned int ucer_seq_hdr_len = 0;
+		asn1_write_element(&p_ucer_seq, &ucer_seq_hdr_len, ASN1_IA5_STRING, (void*)"ucer", -1);
+		asn1_write_element_header(ASN1_OCTET_STRING, ucer_size, &p_ucer_seq, &ucer_seq_hdr_len);
+
+		p = &elem_seq[0];
+		seq_hdr_len = 0;
+		asn1_write_element_header(ASN1_SEQUENCE | ASN1_CONSTRUCTED, ucer_seq_hdr_len + ucer_size, &p, &seq_hdr_len);
+
+		// add size to priv ucer element
+		asn1_write_size(ucer_seq_hdr_len + ucer_size + seq_hdr_len, &p_im4rset, &im4rlen);
+
+		// put it together
+		memcpy(p_im4rset, elem_seq, seq_hdr_len);
+		p_im4rset += seq_hdr_len;
+		im4rlen += seq_hdr_len;
+		memcpy(p_im4rset, ucer_seq, ucer_seq_hdr_len);
+		p_im4rset += ucer_seq_hdr_len;
+		im4rlen += ucer_seq_hdr_len;
+		memcpy(p_im4rset, ucer_data, ucer_size);
+		p_im4rset += ucer_size;
+		im4rlen += ucer_size;
+
+		// now construct IM4R
+
+		/* write inner set */
+		unsigned char inner_set_[8];
+		unsigned char *inner_set = &inner_set_[0];
+		unsigned int inner_set_len = 0;
+		asn1_write_element_header(ASN1_SET | ASN1_CONSTRUCTED, im4rlen, &inner_set, &inner_set_len);
+
+		/* write header values */
+		unsigned char hdrdata_[16];
+		unsigned char *hdrdata = &hdrdata_[0];
+		unsigned int hdrdata_len = 0;
+		asn1_write_element(&hdrdata, &hdrdata_len, ASN1_IA5_STRING, (void*)"IM4R", -1);
+
+		/* write sequence now that we know the entire size */
+		unsigned char seq_[8];
+		unsigned char *seq = &seq_[0];
+		unsigned int seq_len = 0;
+		asn1_write_element_header(ASN1_SEQUENCE | ASN1_CONSTRUCTED, im4rlen + inner_set_len + hdrdata_len, &seq, &seq_len);
+
+		/* write outer cont[1] */
+		unsigned char cont_[8];
+		unsigned char *cont = &cont_[0];
+		unsigned int cont_len = 0;
+		asn1_write_element_header(ASN1_CONTEXT_SPECIFIC | ASN1_CONSTRUCTED | 1, im4rlen + inner_set_len + hdrdata_len + seq_len, &cont, &cont_len);
+
+		// now put everything together
+		additional_data = malloc(im4rlen + inner_set_len + hdrdata_len + seq_len + cont_len);
+		p = additional_data;
+		memcpy(p, cont_, cont_len);
+		p += cont_len;
+		memcpy(p, seq_, seq_len);
+		p += seq_len;
+		memcpy(p, hdrdata_, hdrdata_len);
+		p += hdrdata_len;
+		memcpy(p, inner_set_, inner_set_len);
+		p += inner_set_len;
+		memcpy(p, im4rset, im4rlen);
+		p += im4rlen;
+		additional_size = (unsigned int)(p - additional_data);
+
+		free(im4rset);
+	}
+
 	// create element header for the "IMG4" magic
 	asn1_create_element_header(ASN1_IA5_STRING, IMG4_MAGIC_SIZE, &magic_header, &magic_header_size);
 	// create element header for the blob (ApImg4Ticket)
 	asn1_create_element_header(ASN1_CONTEXT_SPECIFIC|ASN1_CONSTRUCTED, blob_size, &blob_header, &blob_header_size);
 
 	// calculate the size for the final IMG4 file (asn1 sequence)
-	content_size = magic_header_size + IMG4_MAGIC_SIZE + component_size + blob_header_size + blob_size;
+	content_size = magic_header_size + IMG4_MAGIC_SIZE + component_size + blob_header_size + blob_size + additional_size;
 
 	// create element header for the final IMG4 asn1 blob
 	asn1_create_element_header(ASN1_SEQUENCE|ASN1_CONSTRUCTED, content_size, &img4header, &img4header_size);
@@ -457,6 +635,7 @@ int img4_stitch_component(const char* component_name, const unsigned char* compo
 		if (img4header) {
 			free(img4header);
 		}
+		free(additional_data);
 		error("ERROR: out of memory when personalizing IMG4 component %s\n", component_name);
 		return -1;
 	}
@@ -475,6 +654,10 @@ int img4_stitch_component(const char* component_name, const unsigned char* compo
 	p += blob_header_size;
 	memcpy(p, blob, blob_size);
 	p += blob_size;
+	if (additional_size) {
+		memcpy(p, additional_data, additional_size);
+		p += additional_size;
+	}
 
 	*img4_data = outbuf;
 	*img4_size = (p - outbuf);
@@ -488,6 +671,7 @@ int img4_stitch_component(const char* component_name, const unsigned char* compo
 	if (img4header) {
 		free(img4header);
 	}
+	free(additional_data);
 
 	return 0;
 }
diff --git a/src/img4.h b/src/img4.h
index 37dea56..1056fa6 100644
--- a/src/img4.h
+++ b/src/img4.h
@@ -26,7 +26,7 @@
 extern "C" {
 #endif
 
-int img4_stitch_component(const char* component_name, const unsigned char* component_data, unsigned int component_size, const unsigned char* blob, unsigned int blob_size, unsigned char** img4_data, unsigned int *img4_size);
+int img4_stitch_component(const char* component_name, const unsigned char* component_data, unsigned int component_size, plist_t tss_response, unsigned char** img4_data, unsigned int *img4_size);
 int img4_create_local_manifest(plist_t request, plist_t build_identity, plist_t* manifest);
 
 #ifdef __cplusplus