Changed the device type to a structure array for cleaner code and cross state access

7 files changed, 194 insertions, 113 deletions

diff --git a/src/dfu.c b/src/dfu.c
index 1a5a037..1b0e8e7 100644
--- a/src/dfu.c
+++ b/src/dfu.c
@@ -80,7 +80,7 @@ int dfu_enter_recovery(const char* ipsw, plist_t tss) {
 		return -1;
 	}
 
-	idevicerestore_mode = RECOVERY_MODE;
+	idevicerestore_mode = MODE_RECOVERY;
 	irecv_close(dfu);
 	dfu = NULL;
 	return 0;
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index c6a532a..0e77447 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -45,8 +45,8 @@ int idevicerestore_erase = 0;
 int idevicerestore_custom = 0;
 int idevicerestore_verbose = 0;
 int idevicerestore_exclude = 0;
-idevicerestore_mode_t idevicerestore_mode = UNKNOWN_MODE;
-idevicerestore_device_t idevicerestore_device = UNKNOWN_DEVICE;
+int idevicerestore_mode = MODE_UNKNOWN;
+idevicerestore_device_t* idevicerestore_device = NULL;
 
 static struct option long_opts[] = {
 	{ "uuid",    required_argument,  NULL, 'u' },
@@ -79,7 +79,7 @@ int main(int argc, char* argv[]) {
 	char* ipsw = NULL;
 	char* uuid = NULL;
 	uint64_t ecid = 0;
-	while (opt = getopt_long(argc, argv, "vdhceu:", long_opts, &optindex) > 0) {
+	while ((opt = getopt_long(argc, argv, "vdhcexu:", long_opts, &optindex)) > 0) {
 		switch (opt) {
 		case 'h':
 			usage(argc, argv);
@@ -125,6 +125,10 @@ int main(int argc, char* argv[]) {
 		return -1;
 	}
 
+	if(idevicerestore_debug) {
+		idevice_set_debug_level(5);
+	}
+
 	// check which mode the device is currently in so we know where to start
 	idevicerestore_mode = check_mode(uuid);
 	if (idevicerestore_mode < 0) {
@@ -174,7 +178,7 @@ int main(int argc, char* argv[]) {
 	// devices that come after iPod2g require personalized firmwares
 	plist_t tss_request = NULL;
 	plist_t tss = NULL;
-	if (idevicerestore_device > IPOD2G_DEVICE) {
+	if (idevicerestore_device > DEVICE_IPOD2G) {
 		info("Creating TSS request\n");
 		// fetch the device's ECID for the TSS request
 		if (get_ecid(uuid, &ecid) < 0 || ecid == 0) {
@@ -222,7 +226,7 @@ int main(int argc, char* argv[]) {
 	}
 
 	// if the device is in normal mode, place device into recovery mode
-	if (idevicerestore_mode == NORMAL_MODE) {
+	if (idevicerestore_mode == MODE_NORMAL) {
 		info("Entering recovery mode...\n");
 		if (normal_enter_recovery(uuid) < 0) {
 			error("ERROR: Unable to place device into recovery mode\n");
@@ -234,7 +238,7 @@ int main(int argc, char* argv[]) {
 	}
 
 	// if the device is in DFU mode, place device into recovery mode
-	if (idevicerestore_mode == DFU_MODE) {
+	if (idevicerestore_mode == MODE_DFU) {
 		if (dfu_enter_recovery(ipsw, tss) < 0) {
 			error("ERROR: Unable to place device into recovery mode\n");
 			plist_free(buildmanifest);
@@ -245,7 +249,7 @@ int main(int argc, char* argv[]) {
 	}
 
 	// if the device is in recovery mode, place device into restore mode
-	if (idevicerestore_mode == RECOVERY_MODE) {
+	if (idevicerestore_mode == MODE_RECOVERY) {
 		if (recovery_enter_restore(uuid, ipsw, tss) < 0) {
 			error("ERROR: Unable to place device into restore mode\n");
 			plist_free(buildmanifest);
@@ -256,7 +260,7 @@ int main(int argc, char* argv[]) {
 	}
 
 	// device is finally in restore mode, let's do this
-	if (idevicerestore_mode == RESTORE_MODE) {
+	if (idevicerestore_mode == MODE_RESTORE) {
 		info("Restoring device... \n");
 		if (restore_device(uuid, ipsw, tss, filesystem) < 0) {
 			error("ERROR: Unable to restore device\n");
@@ -265,7 +269,7 @@ int main(int argc, char* argv[]) {
 	}
 
 	// device has finished restoring, lets see if we need to activate
-	if (idevicerestore_mode == NORMAL_MODE) {
+	if (idevicerestore_mode == MODE_NORMAL) {
 		info("Checking activation status\n");
 		int activation = activate_check_status(uuid);
 		if (activation < 0) {
@@ -291,52 +295,60 @@ int main(int argc, char* argv[]) {
 }
 
 int check_mode(const char* uuid) {
-	idevicerestore_mode_t mode = UNKNOWN_MODE;
-	if (normal_check_mode(uuid) == 0) {
-		info("Found device in normal mode\n");
-		mode = NORMAL_MODE;
-	}
+	int mode = MODE_UNKNOWN;
 
-	else if (recovery_check_mode() == 0) {
+	if (recovery_check_mode() == 0) {
 		info("Found device in recovery mode\n");
-		mode = RECOVERY_MODE;
+		mode = MODE_RECOVERY;
 	}
 
 	else if (dfu_check_mode() == 0) {
 		info("Found device in DFU mode\n");
-		mode = DFU_MODE;
+		mode = MODE_DFU;
+	}
+
+	else if (normal_check_mode(uuid) == 0) {
+		info("Found device in normal mode\n");
+		mode = MODE_NORMAL;
 	}
 
 	else if (restore_check_mode(uuid) == 0) {
 		info("Found device in restore mode\n");
-		mode = RESTORE_MODE;
+		mode = MODE_RESTORE;
 	}
 
 	return mode;
 }
 
 int check_device(const char* uuid) {
+	int device = DEVICE_UNKNOWN;
 	uint32_t bdid = 0;
 	uint32_t cpid = 0;
-	idevicerestore_device_t device = UNKNOWN_DEVICE;
 
 	switch (idevicerestore_mode) {
-	case NORMAL_MODE:
+	case MODE_RESTORE:
+		device = restore_check_device(uuid);
+		if (device < 0) {
+			device = DEVICE_UNKNOWN;
+		}
+		break;
+
+	case MODE_NORMAL:
 		device = normal_check_device(uuid);
 		if (device < 0) {
-			device = UNKNOWN_DEVICE;
+			device = DEVICE_UNKNOWN;
 		}
 		break;
 
-	case DFU_MODE:
-	case RECOVERY_MODE:
+	case MODE_DFU:
+	case MODE_RECOVERY:
 		if (get_cpid(uuid, &cpid) < 0) {
 			error("ERROR: Unable to get device CPID\n");
 			break;
 		}
 
 		switch (cpid) {
-		case IPHONE2G_CPID:
+		case CPID_IPHONE2G:
 			// iPhone1,1 iPhone1,2 and iPod1,1 all share the same ChipID
 			//   so we need to check the BoardID
 			if (get_bdid(uuid, &bdid) < 0) {
@@ -345,48 +357,48 @@ int check_device(const char* uuid) {
 			}
 
 			switch (bdid) {
-			case IPHONE2G_BDID:
-				device = IPHONE2G_DEVICE;
+			case BDID_IPHONE2G:
+				device = DEVICE_IPHONE2G;
 				break;
 
-			case IPHONE3G_BDID:
-				device = IPHONE3G_DEVICE;
+			case BDID_IPHONE3G:
+				device = DEVICE_IPHONE3G;
 				break;
 
-			case IPOD1G_BDID:
-				device = IPOD1G_DEVICE;
+			case BDID_IPOD1G:
+				device = DEVICE_IPOD1G;
 				break;
 
 			default:
-				device = UNKNOWN_DEVICE;
+				device = DEVICE_UNKNOWN;
 				break;
 			}
 			break;
 
-		case IPHONE3GS_CPID:
-			device = IPHONE3GS_DEVICE;
+		case CPID_IPHONE3GS:
+			device = DEVICE_IPHONE3GS;
 			break;
 
-		case IPOD2G_CPID:
-			device = IPOD2G_DEVICE;
+		case CPID_IPOD2G:
+			device = DEVICE_IPOD2G;
 			break;
 
-		case IPOD3G_CPID:
-			device = IPOD3G_DEVICE;
+		case CPID_IPOD3G:
+			device = DEVICE_IPOD3G;
 			break;
 
-		case IPAD1G_CPID:
-			device = IPAD1G_DEVICE;
+		case CPID_IPAD1G:
+			device = DEVICE_IPAD1G;
 			break;
 
 		default:
-			device = UNKNOWN_DEVICE;
+			device = DEVICE_UNKNOWN;
 			break;
 		}
 		break;
 
 	default:
-		device = UNKNOWN_MODE;
+		device = DEVICE_UNKNOWN;
 		break;
 
 	}
@@ -396,15 +408,15 @@ int check_device(const char* uuid) {
 
 int get_bdid(const char* uuid, uint32_t* bdid) {
 	switch (idevicerestore_mode) {
-	case NORMAL_MODE:
+	case MODE_NORMAL:
 		if (normal_get_bdid(uuid, bdid) < 0) {
 			*bdid = -1;
 			return -1;
 		}
 		break;
 
-	case DFU_MODE:
-	case RECOVERY_MODE:
+	case MODE_DFU:
+	case MODE_RECOVERY:
 		if (recovery_get_bdid(bdid) < 0) {
 			*bdid = -1;
 			return -1;
@@ -421,15 +433,15 @@ int get_bdid(const char* uuid, uint32_t* bdid) {
 
 int get_cpid(const char* uuid, uint32_t* cpid) {
 	switch (idevicerestore_mode) {
-	case NORMAL_MODE:
+	case MODE_NORMAL:
 		if (normal_get_cpid(uuid, cpid) < 0) {
 			*cpid = 0;
 			return -1;
 		}
 		break;
 
-	case DFU_MODE:
-	case RECOVERY_MODE:
+	case MODE_DFU:
+	case MODE_RECOVERY:
 		if (recovery_get_cpid(cpid) < 0) {
 			*cpid = 0;
 			return -1;
@@ -446,15 +458,15 @@ int get_cpid(const char* uuid, uint32_t* cpid) {
 
 int get_ecid(const char* uuid, uint64_t* ecid) {
 	switch (idevicerestore_mode) {
-	case NORMAL_MODE:
+	case MODE_NORMAL:
 		if (normal_get_ecid(uuid, ecid) < 0) {
 			*ecid = 0;
 			return -1;
 		}
 		break;
 
-	case DFU_MODE:
-	case RECOVERY_MODE:
+	case MODE_DFU:
+	case MODE_RECOVERY:
 		if (recovery_get_ecid(ecid) < 0) {
 			*ecid = 0;
 			return -1;
@@ -472,13 +484,14 @@ int get_ecid(const char* uuid, uint64_t* ecid) {
 int extract_buildmanifest(const char* ipsw, plist_t* buildmanifest) {
 	int size = 0;
 	char* data = NULL;
-	if (idevicerestore_device >= IPHONE2G_DEVICE && idevicerestore_device <= IPOD2G_DEVICE) {
+	int device = idevicerestore_device->device_id;
+	if (device >= DEVICE_IPHONE2G && device <= DEVICE_IPOD2G) {
 		// Older devices that don't require personalized firmwares use BuildManifesto.plist
 		if (ipsw_extract_to_memory(ipsw, "BuildManifesto.plist", &data, &size) < 0) {
 			return -1;
 		}
 
-	} else if (idevicerestore_device >= IPHONE3GS_DEVICE && idevicerestore_device <= IPAD1G_DEVICE) {
+	} else if (device >= DEVICE_IPHONE3GS && device <= DEVICE_IPAD1G) {
 		// Whereas newer devices that do require personalized firmwares use BuildManifest.plist
 		if (ipsw_extract_to_memory(ipsw, "BuildManifest.plist", &data, &size) < 0) {
 			return -1;
@@ -602,7 +615,7 @@ int get_signed_component(char* ipsw, plist_t tss, const char* path, char** data,
 		return -1;
 	}
 
-	if (idevicerestore_device > IPOD2G_DEVICE && idevicerestore_custom == 0) {
+	if (idevicerestore_device->device_id > DEVICE_IPOD2G && idevicerestore_custom == 0) {
 		if (img3_replace_signature(img3, component_blob) < 0) {
 			error("ERROR: Unable to replace IMG3 signature\n");
 			free(component_blob);
diff --git a/src/idevicerestore.h b/src/idevicerestore.h
index adea68e..af66892 100644
--- a/src/idevicerestore.h
+++ b/src/idevicerestore.h
@@ -29,60 +29,66 @@
 #define error(...) fprintf(stderr, __VA_ARGS__)
 #define debug(...) if(idevicerestore_debug >= 1) fprintf(stderr, __VA_ARGS__)
 
-#define IPHONE2G_CPID    8900
-#define IPHONE3G_CPID    8900
-#define IPHONE3GS_CPID   8920
-#define IPOD1G_CPID      8900
-#define IPOD2G_CPID      8720
-#define IPOD3G_CPID      8922
-#define IPAD1G_CPID      8930
+#define MODE_UNKNOWN        -1
+#define MODE_DFU             0
+#define MODE_RECOVERY        1
+#define MODE_RESTORE         2
+#define MODE_NORMAL          3
 
-#define IPHONE2G_BDID       0
-#define IPHONE3G_BDID       4
-#define IPHONE3GS_BDID      0
-#define IPOD1G_BDID         2
-#define IPOD2G_BDID         0
-#define IPOD3G_BDID         2
-#define IPAD1G_BDID         2
+#define CPID_UNKNOWN        -1
+#define CPID_IPHONE2G     8900
+#define CPID_IPOD1G       8900
+#define CPID_IPHONE3G     8900
+#define CPID_IPOD2G       8720
+#define CPID_IPHONE3GS    8920
+#define CPID_IPOD3G       8922
+#define CPID_IPAD1G       8930
 
-typedef enum {
-	UNKNOWN_MODE =       -1,
-	DFU_MODE =            0,
-	RECOVERY_MODE =       1,
-	RESTORE_MODE =        2,
-	NORMAL_MODE =         3,
-} idevicerestore_mode_t;
+#define BDID_UNKNOWN        -1
+#define BDID_IPHONE2G        0
+#define BDID_IPOD1G          2
+#define BDID_IPHONE3G        4
+#define BDID_IPOD2G          0
+#define BDID_IPHONE3GS       0
+#define BDID_IPOD3G          2
+#define BDID_IPAD1G          2
 
-typedef enum {
-	UNKNOWN_DEVICE =     -1,
-	IPHONE2G_DEVICE =     0,
-	IPHONE3G_DEVICE =     1,
-	IPOD1G_DEVICE =       2,
-	IPOD2G_DEVICE =       3,
-	IPHONE3GS_DEVICE =    4,
-	IPOD3G_DEVICE =       5,
-	IPAD1G_DEVICE =       6
+#define DEVICE_UNKNOWN      -1
+#define DEVICE_IPHONE2G      0
+#define DEVICE_IPOD1G        1
+#define DEVICE_IPHONE3G      2
+#define DEVICE_IPOD2G        3
+#define DEVICE_IPHONE3GS     4
+#define DEVICE_IPOD3G        5
+#define DEVICE_IPAD1G        6
+
+typedef struct {
+	int device_id;
+	const char* product;
+	const char* model;
+	int board_id;
+	int chip_id;
 } idevicerestore_device_t;
 
-static char* idevicerestore_products[] = {
-	"iPhone1,1",
-	"iPhone1,2",
-	"iPhone2,1",
-	"iPod1,1",
-	"iPod2,1",
-	"iPod3,1",
-	"iPad1,1",
-	NULL
+static idevicerestore_device_t idevicerestore_devices[] = {
+	{  0, "iPhone1,1", "M68AP",  0,  8900 },
+	{  1, "iPod1,1",   "N45AP",  2,  8900 },
+	{  2, "iPhone1,2", "N82AP",  4,  8900 },
+	{  3, "iPod2,1",   "N72AP",  0,  8720 },
+	{  4, "iPhone2,1", "N88AP",  0,  8920 },
+	{  5, "iPod3,1",   "N18AP",  2,  8922 },
+	{  6, "iPad1,1",   "K48AP",  2,  8930 },
+	{ -1,  NULL,        NULL,   -1,    -1 }
 };
 
+extern int idevicerestore_mode;
 extern int idevicerestore_quit;
 extern int idevicerestore_debug;
 extern int idevicerestore_erase;
 extern int idevicerestore_custom;
 extern int idevicerestore_exclude;
 extern int idevicerestore_verbose;
-extern idevicerestore_mode_t idevicerestore_mode;
-extern idevicerestore_device_t idevicerestore_device;
+extern idevicerestore_device_t* idevicerestore_device;
 
 int check_mode(const char* uuid);
 int check_device(const char* uuid);
@@ -97,5 +103,12 @@ int get_shsh_blobs(uint64_t ecid, plist_t build_identity, plist_t* tss);
 int extract_filesystem(const char* ipsw, plist_t buildmanifest, char** filesystem);
 int get_signed_component(char* ipsw, plist_t tss, const char* path, char** data, uint32_t* size);
 
+inline static void debug_plist(plist_t plist) {
+	int size = 0;
+	char* data = NULL;
+	plist_to_xml(plist, &data, &size);
+	debug("%s", data);
+	free(data);
+}
 
 #endif
diff --git a/src/normal.c b/src/normal.c
index ab9216c..b9270d8 100644
--- a/src/normal.c
+++ b/src/normal.c
@@ -62,6 +62,7 @@ int normal_check_mode(const char* uuid) {
 }
 
 int normal_check_device(const char* uuid) {
+	int i = 0;
 	idevice_t device = NULL;
 	char* product_type = NULL;
 	plist_t product_type_node = NULL;
@@ -87,30 +88,25 @@ int normal_check_device(const char* uuid) {
 		return -1;
 	}
 
+	lockdownd_client_free(lockdown);
+	idevice_free(device);
+	lockdown = NULL;
+	device = NULL;
+
 	if (!product_type_node || plist_get_node_type(product_type_node) != PLIST_STRING) {
-		if (product_type_node)
-			plist_free(product_type_node);
-		lockdownd_client_free(lockdown);
-		idevice_free(device);
+		if (product_type_node) plist_free(product_type_node);
 		return -1;
 	}
 	plist_get_string_val(product_type_node, &product_type);
 	plist_free(product_type_node);
 
-	lockdownd_client_free(lockdown);
-	idevice_free(device);
-	lockdown = NULL;
-	device = NULL;
-
-	int i = 0;
-	for (i = 0; idevicerestore_products[i] != NULL; i++) {
-		if (!strcmp(product_type, idevicerestore_products[i])) {
-			idevicerestore_device = i;
+	for (i = 0; idevicerestore_devices[i].product != NULL; i++) {
+		if (!strcmp(product_type, idevicerestore_devices[i].product)) {
 			break;
 		}
 	}
 
-	return idevicerestore_device;
+	return idevicerestore_devices[i].device_id;
 }
 
 int normal_enter_recovery(const char* uuid) {
@@ -166,7 +162,7 @@ int normal_enter_recovery(const char* uuid) {
 		return -1;
 	}
 
-	idevicerestore_mode = RECOVERY_MODE;
+	idevicerestore_mode = MODE_RECOVERY;
 	irecv_close(recovery);
 	recovery = NULL;
 	return 0;
diff --git a/src/recovery.c b/src/recovery.c
index a0a1151..e67b10d 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -132,15 +132,15 @@ int recovery_open_with_timeout(irecv_client_t* client) {
 	int i = 0;
 	int attempts = 10;
 	irecv_client_t recovery = NULL;
-	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+	irecv_error_t recovery_error = IRECV_E_UNKNOWN_ERROR;
 
 	for (i = 1; i <= attempts; i++) {
-		recovery_error = irecv_open(client);
+		recovery_error = irecv_open(&recovery);
 		if (recovery_error == IRECV_E_SUCCESS) {
 			break;
 		}
 
-		if (i == attempts) {
+		if (i >= attempts) {
 			error("ERROR: Unable to connect to device in recovery mode\n");
 			return -1;
 		}
diff --git a/src/restore.c b/src/restore.c
index ce9d41b..46a509a 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -80,6 +80,64 @@ int restore_check_mode(const char* uuid) {
 	return 0;
 }
 
+int restore_check_device(const char* uuid) {
+	int i = 0;
+	char* type = NULL;
+	char* model = NULL;
+	plist_t node = NULL;
+	uint64_t version = 0;
+	idevice_t device = NULL;
+	restored_client_t restore = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	restored_error_t restore_error = RESTORE_E_SUCCESS;
+
+	device_error = idevice_new(&device, uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	restore_error = restored_client_new(device, &restore, "idevicerestore");
+	if (restore_error != RESTORE_E_SUCCESS) {
+		idevice_free(device);
+		return -1;
+	}
+
+	restore_error = restored_query_type(restore, &type, &version);
+	if (restore_error != RESTORE_E_SUCCESS) {
+		restored_client_free(restore);
+		idevice_free(device);
+		return -1;
+	}
+
+	restore_error = restored_get_value(restore, "HardwareModel", &node);
+	if(restore_error != RESTORE_E_SUCCESS) {
+		error("ERROR: Unable to get HardwareModel from restored\n");
+		restored_client_free(restore);
+		idevice_free(device);
+		return -1;
+	}
+
+	restored_client_free(restore);
+	idevice_free(device);
+	restore = NULL;
+	device = NULL;
+
+	if (!node || plist_get_node_type(node) != PLIST_STRING) {
+		error("ERROR: Unable to get HardwareModel information\n");
+		if(node) plist_free(node);
+		return -1;
+	}
+	plist_get_string_val(node, &model);
+
+	for(i = 0; idevicerestore_devices[i].model != NULL; i++) {
+		if(!strcasecmp(model, idevicerestore_devices[i].model)) {
+			break;
+		}
+	}
+
+	return idevicerestore_devices[i].device_id;
+}
+
 void restore_device_callback(const idevice_event_t* event, void* user_data) {
 	if (event->event == IDEVICE_DEVICE_ADD) {
 		restore_device_connected = 1;
diff --git a/src/restore.h b/src/restore.h
index 99ba80b..6614355 100644
--- a/src/restore.h
+++ b/src/restore.h
@@ -27,6 +27,7 @@
 #include <libimobiledevice/libimobiledevice.h>
 
 int restore_check_mode(const char* uuid);
+int restore_check_device(const char* uuid);
 const char* restore_progress_string(unsigned int operation);
 void restore_close(idevice_t device, restored_client_t restore);
 int restore_handle_status_msg(restored_client_t client, plist_t msg);