Merge branch 'martin'

22 files changed, 2686 insertions, 1016 deletions

diff --git a/TODO b/TODO
index 80e336c..efd9e9b 100644
--- a/TODO
+++ b/TODO
@@ -1,19 +1,24 @@
-x Unzip IPSW
-x Parse BuildManifest.plist
-x Build TSS request
-x Submit TSS request to apple
-X Stitch new signatures onto firmware
-X Connect to iBoot
-X Upload iBEC and execute go
-X Reconnect to iBEC
-X Upload devicetree and execute devicetree
-X Upload ramdisk and execute ramdisk
-X Upload kernelcache and execute bootx
-X Send QueryType request
-X Send StartRestore request
-X Send source verification packets
-X Send filesystem payload
-X Send kernelcache payload
-X Send NOR firmware payload
-* Fetch activation ticket from apple
-* Send activation ticket to device
\ No newline at end of file
+TODO List
+------------------------------------------------
+
+o) Instead of assumeing upgrade or restore, the program should query all avaiable build identitiys and find which ones are valid, if more then one is valid then display a menu asking whether to resore or upgrade. (This would also benifit users by cacheing both upgrade and restore SHSH blobs to sauriks server so users aren't stuck with only being able to upgrade or restore.) (in progress...)
+o) Also img3 should be updated to be able to fetch tags from each image such as the VERS tag to be able to compare with the current iBoot version. 
+o) Need to add 8900 file support into img3.c. 
+o) Cacheing TSS and activation info locally so the program doesn't need to query apple/cydia servers everytime would be nice. 
+o) Need to add support to handle and display status messages correctly. 
+o) Writing restore logs to an external file for debugging would be a great help. 
+o) Would be nice to add options to keep some connections open during similar operations rather then reopening and closing them each time. (in progress...)
+o) Extracting and signing all firmware files at one time rather then extracting and signing each one as we need it could speed things up as well.
+o) Double check restore_send_kernelcache()
+o) Clean up and refactor restore_send_nor()
+o) Finishing cleaning up asr_perform_validation
+o) Double check asr_send_payload()
+o) Cleanup dfu.c/dfu.h
+o) Cleanup recovery.c/recovery.h
+o) Cleanup normal.c/normal.h
+o) Cleanup restore.c/restore.h
+o) Cleanup tss.c/tss.h
+o) Cleanup img3.c/img3.h
+o) Cleanup ipsw.c/ipsw.h
+o) Make ipsw.c a proxy to tss.c and img3.c
+o) Make all components only have access to their own context instead of a single global application context (less coupling)
\ No newline at end of file
diff --git a/src/Makefile.am b/src/Makefile.am
index 6840a0c..a4c8fc3 100644
--- a/src/Makefile.am
+++ b/src/Makefile.am
@@ -18,6 +18,6 @@ AM_LDFLAGS =\
 
 bin_PROGRAMS = idevicerestore
 
-idevicerestore_SOURCES = idevicerestore.c dfu.c tss.c img3.c ipsw.c normal.c restore.c recovery.c
+idevicerestore_SOURCES = idevicerestore.c common.c tss.c img3.c ipsw.c normal.c dfu.c recovery.c restore.c asr.c
 idevicerestore_CFLAGS = $(AM_CFLAGS)
-idevicerestore_LDFLAGS = $(AM_LDFLAGS)
\ No newline at end of file
+idevicerestore_LDFLAGS = $(AM_LDFLAGS)
diff --git a/src/asr.c b/src/asr.c
new file mode 100644
index 0000000..2331660
--- /dev/null
+++ b/src/asr.c
@@ -0,0 +1,292 @@
+/*
+ * asr.h
+ * Functions for handling asr connections
+ *
+ * Copyright (c) 2010 Joshua Hill. All Rights Reserved.
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <libimobiledevice/libimobiledevice.h>
+
+#include "asr.h"
+#include "idevicerestore.h"
+
+#define ASR_PORT 12345
+#define ASR_BUFFER_SIZE 65536
+
+int asr_open_with_timeout(idevice_t device, idevice_connection_t* asr) {
+	int i = 0;
+	int attempts = 10;
+	idevice_connection_t connection = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+
+	*asr = NULL;
+
+	if (device == NULL) {
+		return -1;
+	}
+
+	debug("Connecting to ASR\n");
+	for (i = 1; i <= attempts; i++) {
+		device_error = idevice_connect(device, ASR_PORT, &connection);
+		if (device_error == IDEVICE_E_SUCCESS) {
+			break;
+		}
+
+		if (i >= attempts) {
+			error("ERROR: Unable to connect to ASR client\n");
+			return -1;
+		}
+
+		sleep(2);
+		debug("Retrying connection...\n");
+	}
+
+	*asr = connection;
+	return 0;
+}
+
+int asr_receive(idevice_connection_t asr, plist_t* data) {
+	uint32_t size = 0;
+	char* buffer = NULL;
+	plist_t request = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+
+	*data = NULL;
+
+	buffer = (char*) malloc(ASR_BUFFER_SIZE);
+	if (buffer == NULL) {
+		error("ERROR: Unable to allocate memory for ASR receive buffer\n");
+		return -1;
+	}
+	memset(buffer, '\0', ASR_BUFFER_SIZE);
+
+	device_error = idevice_connection_receive(asr, buffer, ASR_BUFFER_SIZE, &size);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		error("ERROR: Unable to receive data from ASR\n");
+		free(buffer);
+		return -1;
+	}
+	plist_from_xml(buffer, size, &request);
+
+	*data = request;
+
+	debug("Received %d bytes:\n%s\n", size, buffer);
+	free(buffer);
+	return 0;
+}
+
+int asr_send(idevice_connection_t asr, plist_t* data) {
+	uint32_t size = 0;
+	char* buffer = NULL;
+
+	plist_to_xml(data, &buffer, &size);
+	if (asr_send_buffer(asr, buffer, size) < 0) {
+		error("ERROR: Unable to send plist to ASR\n");
+		free(buffer);
+		return -1;
+	}
+
+	debug("Sent %d bytes:\n", size);
+	debug_plist(data);
+	free(buffer);
+	return 0;
+}
+
+int asr_send_buffer(idevice_connection_t asr, const char* data, uint32_t size) {
+	uint32_t bytes = 0;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+
+	device_error = idevice_connection_send(asr, data, size, &bytes);
+	if (device_error != IDEVICE_E_SUCCESS || bytes != size) {
+		error("ERROR: Unable to send data to ASR\n");
+		return -1;
+	}
+
+	return 0;
+}
+
+void asr_close(idevice_connection_t asr) {
+	if (asr != NULL) {
+		idevice_disconnect(asr);
+		asr = NULL;
+	}
+}
+
+int asr_perform_validation(idevice_connection_t asr, const char* filesystem) {
+	FILE* file = NULL;
+	uint64_t length = 0;
+	char* command = NULL;
+	plist_t node = NULL;
+	plist_t packet = NULL;
+	plist_t packet_info = NULL;
+	plist_t payload_info = NULL;
+
+	file = fopen(filesystem, "rb");
+	if (file == NULL) {
+		return -1;
+	}
+
+	fseek(file, 0, SEEK_END);
+	length = ftell(file);
+	fseek(file, 0, SEEK_SET);
+
+	payload_info = plist_new_dict();
+	plist_dict_insert_item(payload_info, "Port", plist_new_uint(1));
+	plist_dict_insert_item(payload_info, "Size", plist_new_uint(length));
+
+	packet_info = plist_new_dict();
+	plist_dict_insert_item(packet_info, "FEC Slice Stride", plist_new_uint(40));
+	plist_dict_insert_item(packet_info, "Packet Payload Size", plist_new_uint(1450));
+	plist_dict_insert_item(packet_info, "Packets Per FEC", plist_new_uint(25));
+	plist_dict_insert_item(packet_info, "Payload", payload_info);
+	plist_dict_insert_item(packet_info, "Stream ID", plist_new_uint(1));
+	plist_dict_insert_item(packet_info, "Version", plist_new_uint(1));
+
+	if (asr_send(asr, packet_info)) {
+		error("ERROR: Unable to sent packet information to ASR\n");
+		plist_free(packet_info);
+		return -1;
+	}
+	plist_free(packet_info);
+
+	while (1) {
+		if (asr_receive(asr, &packet) < 0) {
+			error("ERROR: Unable to receive validation packet\n");
+			return -1;
+		}
+
+		node = plist_dict_get_item(packet, "Command");
+		if (!node || plist_get_node_type(node) != PLIST_STRING) {
+			error("ERROR: Unable to find command node in validation request\n");
+			return -1;
+		}
+		plist_get_string_val(node, &command);
+
+		if (!strcmp(command, "OOBData")) {
+			asr_handle_oob_data_request(asr, packet, file);
+
+
+			plist_free(packet);
+
+		} else if(!strcmp(command, "Payload")) {
+			plist_free(packet);
+			break;
+
+		} else {
+			error("ERROR: Unknown command received from ASR\n");
+			plist_free(packet);
+			return -1;
+		}
+	}
+
+	return 0;
+}
+
+int asr_handle_oob_data_request(idevice_connection_t asr, plist_t packet, FILE* file) {
+	char* oob_data = NULL;
+	uint64_t oob_offset = 0;
+	uint64_t oob_length = 0;
+	plist_t oob_length_node = NULL;
+	plist_t oob_offset_node = NULL;
+
+	oob_length_node = plist_dict_get_item(packet, "OOB Length");
+	if (!oob_length_node || PLIST_UINT != plist_get_node_type(oob_length_node)) {
+		error("ERROR: Unable to find OOB data length\n");
+		return -1;
+	}
+	plist_get_uint_val(oob_length_node, &oob_length);
+
+	oob_offset_node = plist_dict_get_item(packet, "OOB Offset");
+	if (!oob_offset_node || PLIST_UINT != plist_get_node_type(oob_offset_node)) {
+		error("ERROR: Unable to find OOB data offset\n");
+		return -1;
+	}
+	plist_get_uint_val(oob_offset_node, &oob_offset);
+
+	oob_data = (char*) malloc(oob_length);
+	if (oob_data == NULL) {
+		error("ERROR: Out of memory\n");
+		plist_free(packet);
+		return -1;
+	}
+
+	fseek(file, oob_offset, SEEK_SET);
+	if (fread(oob_data, 1, oob_length, file) != oob_length) {
+		error("ERROR: Unable to read OOB data from filesystem offset\n");
+		plist_free(packet);
+		free(oob_data);
+		return -1;
+	}
+
+	if (asr_send_buffer(asr, oob_data, oob_length) < 0) {
+		error("ERROR: Unable to send OOB data to ASR\n");
+		plist_free(packet);
+		free(oob_data);
+		return -1;
+	}
+	free(oob_data);
+	return 0;
+}
+
+int asr_send_payload(idevice_connection_t asr, const char* filesystem) {
+	int i = 0;
+	char data[1450];
+	FILE* file = NULL;
+	uint32_t bytes = 0;
+	uint32_t count = 0;
+	uint32_t length = 0;
+	double progress = 0;
+
+	file = fopen(filesystem, "rb");
+	if (file == NULL) {
+		return -1;
+	}
+
+	fseek(file, 0, SEEK_END);
+	length = ftell(file);
+	fseek(file, 0, SEEK_SET);
+
+	for(i = length; i > 0; i -= 1450) {
+		int size = 1450;
+		if (i < 1450) {
+			size = i;
+		}
+
+		if (fread(data, 1, size, file) != (unsigned int) size) {
+			error("Error reading filesystem\n");
+			fclose(file);
+			return -1;
+		}
+
+		if (asr_send_buffer(asr, data, size) < 0) {
+			error("ERROR: Unable to send filesystem payload\n");
+			fclose(file);
+			return -1;
+		}
+
+		bytes += size;
+		progress = ((double) bytes/ (double) length) * 100.0;
+		print_progress_bar(progress);
+
+	}
+
+	fclose(file);
+	return 0;
+}
diff --git a/src/asr.h b/src/asr.h
new file mode 100644
index 0000000..1423496
--- /dev/null
+++ b/src/asr.h
@@ -0,0 +1,45 @@
+/*
+ * asr.h
+ * Functions for handling asr connections
+ *
+ * Copyright (c) 2010 Joshua Hill. All Rights Reserved.
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ */
+
+#ifndef IDEVICERESTORE_ASR_H
+#define IDEVICERESTORE_ASR_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <libimobiledevice/libimobiledevice.h>
+
+int asr_open_with_timeout(idevice_t device, idevice_connection_t* asr);
+int asr_send(idevice_connection_t asr, plist_t* data);
+int asr_receive(idevice_connection_t asr, plist_t* data);
+int asr_send_buffer(idevice_connection_t asr, const char* data, uint32_t size);
+void asr_close(idevice_connection_t asr);
+int asr_perform_validation(idevice_connection_t asr, const char* filesystem);
+int asr_send_payload(idevice_connection_t asr, const char* filesystem);
+int asr_handle_oob_data_request(idevice_connection_t asr, plist_t packet, FILE* file);
+
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif
diff --git a/src/common.c b/src/common.c
new file mode 100644
index 0000000..8e76697
--- /dev/null
+++ b/src/common.c
@@ -0,0 +1,110 @@
+/*
+ * common.c
+ * Misc functions used in idevicerestore
+ *
+ * Copyright (c) 2010 Joshua Hill. All Rights Reserved.
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ */
+
+#include <stdio.h>
+
+#include "common.h"
+
+int idevicerestore_debug = 0;
+
+int write_file(const char* filename, const void* data, size_t size) {
+	size_t bytes = 0;
+	FILE* file = NULL;
+
+	debug("Writing data to %s\n", filename);
+	file = fopen(filename, "wb");
+	if (file == NULL) {
+		error("read_file: Unable to open file %s\n", filename);
+		return -1;
+	}
+
+	bytes = fwrite(data, 1, size, file);
+	fclose(file);
+
+	if (bytes != size) {
+		error("ERROR: Unable to write entire file: %s: %d of %d\n", filename, bytes, size);
+		return -1;
+	}
+
+	return size;
+}
+
+int read_file(const char* filename, void** data, size_t* size) {
+	size_t bytes = 0;
+	size_t length = 0;
+	FILE* file = NULL;
+	char* buffer = NULL;
+	debug("Reading data from %s\n", filename);
+
+	*size = 0;
+	*data = NULL;
+
+	file = fopen(filename, "rb");
+	if (file == NULL) {
+		error("read_file: File %s not found\n", filename);
+		return -1;
+	}
+
+	fseek(file, 0, SEEK_END);
+	length = ftell(file);
+	rewind(file);
+
+	buffer = (char*) malloc(length);
+	if (buffer == NULL) {
+		error("ERROR: Out of memory\n");
+		fclose(file);
+		return -1;
+	}
+	bytes = fread(buffer, 1, length, file);
+	fclose(file);
+
+	if (bytes != length) {
+		error("ERROR: Unable to read entire file\n");
+		free(buffer);
+		return -1;
+	}
+
+	*size = length;
+	*data = buffer;
+	return 0;
+}
+
+void debug_plist(plist_t plist) {
+	int size = 0;
+	char* data = NULL;
+	plist_to_xml(plist, &data, &size);
+	info("%s", data);
+	free(data);
+}
+
+void print_progress_bar(double progress) {
+	int i = 0;
+	if(progress < 0) return;
+	if(progress > 100) progress = 100;
+	info("\r[");
+	for(i = 0; i < 50; i++) {
+		if(i < progress / 2) info("=");
+		else info(" ");
+	}
+	info("] %3.1f%%", progress);
+	if(progress == 100) info("\n");
+	fflush(stdout);
+}
diff --git a/src/common.h b/src/common.h
new file mode 100644
index 0000000..f8c0ae0
--- /dev/null
+++ b/src/common.h
@@ -0,0 +1,154 @@
+/*
+ * common.h
+ * Misc functions used in idevicerestore
+ *
+ * Copyright (c) 2010 Joshua Hill. All Rights Reserved.
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
+ */
+
+#ifndef IDEVICERESTORE_COMMON_H
+#define IDEVICERESTORE_COMMON_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <plist/plist.h>
+
+#define info(...) printf(__VA_ARGS__)
+#define error(...) fprintf(stderr, __VA_ARGS__)
+#define debug(...) if(idevicerestore_debug) fprintf(stderr, __VA_ARGS__)
+
+#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
+#define CPID_IPHONE4      8930
+
+#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
+#define BDID_IPHONE4         0
+
+#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
+#define DEVICE_IPHONE4       7
+
+#define MODE_UNKNOWN        -1
+#define MODE_DFU             0
+#define MODE_RECOVERY        1
+#define MODE_RESTORE         2
+#define MODE_NORMAL          3
+
+#define FLAG_QUIT            1
+#define FLAG_DEBUG           2
+#define FLAG_ERASE           4
+#define FLAG_CUSTOM          8
+#define FLAG_EXCLUDE        16
+
+struct dfu_client_t;
+struct normal_client_t;
+struct restore_clien_t;
+struct recovery_client_t;
+
+struct idevicerestore_mode_t {
+	int index;
+	const char* string;
+};
+
+struct idevicerestore_entry_t {
+	char* name;
+	char* path;
+	char* filename;
+	char* blob_data;
+	uint32_t blob_size;
+	struct idevicerestore_entry* next;
+	struct idevicerestore_entry* prev;
+};
+
+struct idevicerestore_device_t {
+	int index;
+	const char* product;
+	const char* model;
+	uint32_t board_id;
+	uint32_t chip_id;
+};
+
+struct idevicerestore_client_t {
+	int flags;
+	plist_t tss;
+	uint64_t ecid;
+	const char* uuid;
+	const char* ipsw;
+	const char* filesystem;
+	struct dfu_client_t* dfu;
+	struct normal_client_t* normal;
+	struct restore_client_t* restore;
+	struct recovery_client_t* recovery;
+	struct idevicerestore_device_t* device;
+	struct idevicerestore_entry_t** entries;
+	struct idevicerestore_mode_t* mode;
+};
+
+static struct idevicerestore_mode_t idevicerestore_modes[] = {
+	{  0, "DFU"      },
+	{  1, "Recovery" },
+	{  2, "Restore"  },
+	{  3, "Normal"   },
+	{ -1,  NULL      }
+};
+
+static struct 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 },
+	{  7, "iPhone3,1", "N90AP",  0,  8930 },
+	{ -1,  NULL,        NULL,   -1,    -1 }
+};
+
+extern int idevicerestore_debug;
+
+void debug_plist(plist_t plist);
+void print_progress_bar(double progress);
+int read_file(const char* filename, void** data, size_t* size);
+int write_file(const char* filename, const void* data, size_t size);
+
+extern struct idevicerestore_client_t* idevicerestore;
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif
diff --git a/src/dfu.c b/src/dfu.c
index 5e13f38..0b3ad51 100644
--- a/src/dfu.c
+++ b/src/dfu.c
@@ -19,10 +19,133 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <libirecovery.h>
 
 #include "dfu.h"
+#include "recovery.h"
+#include "idevicerestore.h"
 
-int dfu_get_ecid(uint64_t* ecid) {
+int dfu_progress_callback(irecv_client_t client, const irecv_event_t* event) {
+	if (event->type == IRECV_PROGRESS) {
+		print_progress_bar(event->progress);
+	}
 	return 0;
 }
+
+int dfu_client_new(struct idevicerestore_client_t* client, uint32_t timeout) {
+	struct dfu_client_t* dfu = NULL;
+	if(client == NULL) {
+		return -1;
+	}
+
+	if(client->dfu) {
+		dfu_client_free(client);
+	}
+
+	dfu = (struct dfu_client_t*) malloc(sizeof(struct dfu_client_t));
+	if (dfu == NULL) {
+		error("ERROR: Out of memory\n");
+		return -1;
+	}
+
+	if (dfu_open_with_timeout(dfu, timeout) < 0) {
+		dfu_client_free(client);
+		return -1;
+	}
+
+	if(dfu->client->mode != kDfuMode) {
+		dfu_client_free(client);
+		return -1;
+	}
+
+	client->dfu = dfu;
+	return 0;
+}
+
+void dfu_client_free(struct idevicerestore_client_t* client) {
+	struct dfu_client_t* dfu = NULL;
+	if(client != NULL) {
+		dfu = client->dfu;
+		if (dfu != NULL) {
+			if(dfu->client != NULL) {
+				irecv_close(dfu->client);
+				dfu->client = NULL;
+			}
+			free(dfu);
+		}
+		client->dfu = NULL;
+	}
+}
+
+int dfu_open_with_timeout(struct idevicerestore_client_t* client, uint32_t timeout) {
+	int i = 0;
+	irecv_client_t recovery = NULL;
+	irecv_error_t recovery_error = IRECV_E_UNKNOWN_ERROR;
+
+	for (i = 1; i <= timeout; i++) {
+		recovery_error = irecv_open(&recovery);
+		if (recovery_error == IRECV_E_SUCCESS) {
+			break;
+		}
+
+		if (i == timeout) {
+			error("ERROR: Unable to connect to device in DFU mode\n");
+			return -1;
+		}
+
+		sleep(1);
+		debug("Retrying connection...\n");
+	}
+
+	irecv_event_subscribe(recovery, IRECV_PROGRESS, &dfu_progress_callback, NULL);
+	client->dfu->client = recovery;
+	return 0;
+}
+
+int dfu_check_mode() {
+	return -1;
+}
+
+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)
+			irecv_close(dfu);
+		return -1;
+	}
+
+	if (recovery_send_component(client, build_identity, component) < 0) {
+		error("ERROR: Unable to send %s to device\n", component);
+		irecv_close(dfu);
+		return -1;
+	}
+
+	dfu_error = irecv_reset(client->dfu->client);
+	if (dfu_error != IRECV_E_SUCCESS) {
+		error("ERROR: Unable to reset device\n");
+		irecv_close(dfu);
+		return -1;
+	}
+	irecv_close(client->dfu->client);
+	client->dfu->client = NULL;
+
+	// Reconnect to device, but this time make sure we're not still in DFU mode
+	if (recovery_open_with_timeout(client) < 0 || client->mode->index != kDfuMode) {
+		error("ERROR: Unable to connect to recovery device\n");
+		if (client->dfu->client)
+			irecv_close(client->dfu->client);
+		return -1;
+	}
+
+	client->mode = &idevicerestore_modes[MODE_RECOVERY];
+	irecv_close(client->dfu->client);
+	client->dfu->client = NULL;
+	return 0;
+}
+
diff --git a/src/dfu.h b/src/dfu.h
index ef9d911..259cb6a 100644
--- a/src/dfu.h
+++ b/src/dfu.h
@@ -19,11 +19,29 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef DFU_H
-#define DFU_H
+#ifndef IDEVICERESTORE_DFU_H
+#define IDEVICERESTORE_DFU_H
 
-#include <stdint.h>
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <libirecovery.h>
+#include "common.h"
+
+struct dfu_client_t {
+	irecv_client_t client;
+	const char* ipsw;
+	plist_t tss;
+};
 
-int dfu_get_ecid(uint64_t* ecid);
+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, plist_t build_identity);
+
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif
diff --git a/src/idevicerestore.c b/src/idevicerestore.c
index aaff4d6..7982ed7 100644
--- a/src/idevicerestore.c
+++ b/src/idevicerestore.c
@@ -23,61 +23,78 @@
 #include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
+#include <getopt.h>
 #include <plist/plist.h>
-#include <libirecovery.h>
-#include <libimobiledevice/restore.h>
-#include <libimobiledevice/lockdown.h>
-#include <libimobiledevice/libimobiledevice.h>
 
 #include "dfu.h"
 #include "tss.h"
 #include "img3.h"
 #include "ipsw.h"
+#include "common.h"
 #include "normal.h"
 #include "restore.h"
 #include "recovery.h"
 #include "idevicerestore.h"
 
-#define UNKNOWN_MODE   0
-#define DFU_MODE       1
-#define NORMAL_MODE    2
-#define RECOVERY_MODE  3
-#define RESTORE_MODE   4
-
-int idevicerestore_debug = 0;
-static int idevicerestore_mode = 0;
-static int idevicerestore_quit = 0;
-static int idevicerestore_custom = 0;
-
-void usage(int argc, char* argv[]);
-int write_file(const char* filename, char* data, int size);
-int get_tss_data_by_name(plist_t tss, const char* entry, char** path, char** blob);
-int get_tss_data_by_path(plist_t tss, const char* path, char** name, char** blob);
-void device_callback(const idevice_event_t* event, void *user_data);
-int get_signed_component_by_name(char* ipsw, plist_t tss, char* component, char** pdata, int* psize);
-int get_signed_component_by_path(char* ipsw, plist_t tss, char* path, char** pdata, int* psize);
+static struct option longopts[] = {
+	{ "uuid",    required_argument, NULL, 'u' },
+	{ "debug",   no_argument,       NULL, 'd' },
+	{ "help",    no_argument,       NULL, 'h' },
+	{ "erase",   no_argument,       NULL, 'e' },
+	{ "custom",  no_argument,       NULL, 'c' },
+	{ "exclude", no_argument,       NULL, 'x' },
+	{ NULL, 0, NULL, 0 }
+};
+
+void usage(int argc, char* argv[]) {
+	char* name = strrchr(argv[0], '/');
+	printf("Usage: %s [OPTIONS] FILE\n", (name ? name + 1 : argv[0]));
+	printf("Restore/upgrade IPSW firmware FILE to an iPhone/iPod Touch.\n");
+	printf("  -u, --uuid UUID\ttarget specific device by its 40-digit device UUID\n");
+	printf("  -d, --debug\t\tenable communication debugging\n");
+	printf("  -h, --help\t\tprints usage information\n");
+	printf("  -e, --erase\t\tperform a full restore, erasing all data\n");
+	printf("  -c, --custom\t\trestore with a custom firmware\n");
+	printf("  -x, --exclude\t\texclude nor/baseband upgrade\n");
+	printf("\n");
+}
 
 int main(int argc, char* argv[]) {
 	int opt = 0;
+	int optindex = 0;
 	char* ipsw = NULL;
 	char* uuid = NULL;
 	uint64_t ecid = 0;
-	while ((opt = getopt(argc, argv, "vdhcu:")) > 0) {
+
+	// create an instance of our context
+	struct idevicerestore_client_t* client = (struct idevicerestore_client_t*) malloc(sizeof(struct idevicerestore_client_t));
+	if (client == NULL) {
+		error("ERROR: Out of memory\n");
+		return -1;
+	}
+	memset(client, '\0', sizeof(struct idevicerestore_client_t));
+
+	while ((opt = getopt_long(argc, argv, "dhcexu:", longopts, &optindex)) > 0) {
 		switch (opt) {
 		case 'h':
 			usage(argc, argv);
+			return 0;
+
+		case 'd':
+			client->flags &= FLAG_DEBUG;
+			idevicerestore_debug = 1;
 			break;
 
-		case 'v':
-			idevicerestore_debug += 1;
+		case 'e':
+			client->flags &= FLAG_ERASE;
 			break;
 
 		case 'c':
-			idevicerestore_custom = 1;
+			client->flags &= FLAG_CUSTOM;
 			break;
 
-		case 'd':
-			idevicerestore_debug = 3;
+		case 'x':
+			client->flags &= FLAG_EXCLUDE;
 			break;
 
 		case 'u':
@@ -86,604 +103,556 @@ int main(int argc, char* argv[]) {
 
 		default:
 			usage(argc, argv);
-			break;
+			return -1;
 		}
 	}
 
-	argc -= optind;
-	argv += optind;
+	if ((argc-optind) == 1) {
+		argc -= optind;
+		argv += optind;
 
-	if (argc == 1)
 		ipsw = argv[0];
-
-	if (ipsw == NULL) {
-		error("ERROR: Please supply an IPSW\n");
+	} else {
+		usage(argc, argv);
 		return -1;
 	}
 
-	idevice_t device = NULL;
-	irecv_client_t recovery = NULL;
-	lockdownd_client_t lockdown = NULL;
-	irecv_error_t recovery_error = IRECV_E_SUCCESS;
-	idevice_error_t device_error = IDEVICE_E_SUCCESS;
-	lockdownd_error_t lockdown_error = LOCKDOWN_E_SUCCESS;
-
-	/* determine recovery or normal mode */
-	info("Checking for device in normal mode...\n");
-	device_error = idevice_new(&device, uuid);
-	if (device_error != IDEVICE_E_SUCCESS) {
-		info("Checking for the device in recovery mode...\n");
-		recovery_error = irecv_open(&recovery);
-		if (recovery_error != IRECV_E_SUCCESS) {
-			error("ERROR: Unable to find device, is it plugged in?\n");
-			return -1;
-		}
-		info("Found device in recovery mode\n");
-		idevicerestore_mode = RECOVERY_MODE;
-
-	} else {
-		info("Found device in normal mode\n");
-		idevicerestore_mode = NORMAL_MODE;
+	if (client->flags & FLAG_DEBUG) {
+		idevice_set_debug_level(1);
+		irecv_set_debug_level(1);
 	}
 
-	/* retrieve ECID */
-	if (idevicerestore_mode == NORMAL_MODE) {
-		lockdown_error = lockdownd_client_new_with_handshake(device, &lockdown, "idevicerestore");
-		if (lockdown_error != LOCKDOWN_E_SUCCESS) {
-			error("ERROR: Unable to connect to lockdownd\n");
-			idevice_free(device);
-			return -1;
-		}
+	client->uuid = uuid;
+	client->ipsw = ipsw;
 
-		plist_t unique_chip_node = NULL;
-		lockdown_error = lockdownd_get_value(lockdown, NULL, "UniqueChipID", &unique_chip_node);
-		if (lockdown_error != LOCKDOWN_E_SUCCESS) {
-			error("ERROR: Unable to get UniqueChipID from lockdownd\n");
-			lockdownd_client_free(lockdown);
-			idevice_free(device);
-			return -1;
-		}
-
-		if (!unique_chip_node || plist_get_node_type(unique_chip_node) != PLIST_UINT) {
-			error("ERROR: Unable to get ECID\n");
-			lockdownd_client_free(lockdown);
-			idevice_free(device);
-			return -1;
-		}
-
-		plist_get_uint_val(unique_chip_node, &ecid);
-		lockdownd_client_free(lockdown);
-		plist_free(unique_chip_node);
-		idevice_free(device);
-		lockdown = NULL;
-		device = NULL;
-
-	} else if (idevicerestore_mode == RECOVERY_MODE) {
-		recovery_error = irecv_get_ecid(recovery, &ecid);
-		if (recovery_error != IRECV_E_SUCCESS) {
-			error("ERROR: Unable to get device ECID\n");
-			irecv_close(recovery);
-			return -1;
-		}
-		irecv_close(recovery);
-		recovery = NULL;
-	}
-
-	if (ecid != 0) {
-		info("Found ECID %llu\n", ecid);
-	} else {
-		error("Unable to find device ECID\n");
+	// check which mode the device is currently in so we know where to start
+	if (check_mode(client) < 0 || client->mode->index == MODE_UNKNOWN) {
+		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);
 
-	/* parse buildmanifest */
-	int buildmanifest_size = 0;
-	char* buildmanifest_data = NULL;
-	info("Extracting BuildManifest.plist from IPSW\n");
-	if (ipsw_extract_to_memory(ipsw, "BuildManifest.plist", &buildmanifest_data, &buildmanifest_size) < 0) {
-		error("ERROR: Unable to extract BuildManifest.plist IPSW\n");
+	// discover the device type
+	if (check_device(client) < 0 || client->device->index == DEVICE_UNKNOWN) {
+		error("ERROR: Unable to discover device type\n");
 		return -1;
 	}
+	info("Identified device as %s\n", client->device->product);
 
-	plist_t manifest = NULL;
-	plist_from_xml(buildmanifest_data, buildmanifest_size, &manifest);
-
-	info("Creating TSS request\n");
-	plist_t tss_request = tss_create_request(manifest, ecid);
-	if (tss_request == NULL) {
-		error("ERROR: Unable to create TSS request\n");
-		plist_free(manifest);
-		return -1;
+	if (client->mode->index == MODE_RESTORE) {
+		if (restore_reboot(client) < 0) {
+			error("ERROR: Unable to exit restore mode\n");
+			return -1;
+		}
 	}
-	plist_free(manifest);
 
-	info("Sending TSS request\n");
-	plist_t tss_response = tss_send_request(tss_request);
-	if (tss_response == NULL) {
-		error("ERROR: Unable to get response from TSS server\n");
-		plist_free(tss_request);
+	// extract buildmanifest
+	plist_t buildmanifest = NULL;
+	info("Extracting BuildManifest from IPSW\n");
+	if (extract_buildmanifest(client, ipsw, &buildmanifest) < 0) {
+		error("ERROR: Unable to extract BuildManifest from %s\n", ipsw);
 		return -1;
 	}
-	info("Got TSS response\n");
 
-	// Get name of filesystem DMG in IPSW
-	char* filesystem = NULL;
-	plist_t filesystem_node = plist_dict_get_item(tss_request, "OS");
-	if (!filesystem_node || plist_get_node_type(filesystem_node) != PLIST_DICT) {
-		error("ERROR: Unable to find filesystem node\n");
-		plist_free(tss_request);
-		return -1;
+	// devices are listed in order from oldest to newest
+	// so we'll need their ECID
+	if (client->device->index > DEVICE_IPOD2G) {
+		debug("Creating TSS request\n");
+		// fetch the device's ECID for the TSS request
+		if (get_ecid(client, &client->ecid) < 0) {
+			error("ERROR: Unable to find device ECID\n");
+			return -1;
+		}
+		debug("Found ECID %llu\n", client->ecid);
 	}
 
-	plist_t filesystem_info_node = plist_dict_get_item(filesystem_node, "Info");
-	if (!filesystem_info_node || plist_get_node_type(filesystem_info_node) != PLIST_DICT) {
-		error("ERROR: Unable to find filesystem info node\n");
-		plist_free(tss_request);
-		return -1;
+	// choose whether this is an upgrade or a restore (default to upgrade)
+	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 any build identities\n");
+			plist_free(buildmanifest);
+			return -1;
+		}
+	} else {
+		// loop through all build identities in the build manifest
+		// and list the valid ones
+		int i = 0;
+		int valid_builds = 0;
+		int build_count = get_build_count(buildmanifest);
+		for (i = 0; i < build_count; i++) {
+			build_identity = get_build_identity(client, buildmanifest, i);
+			valid_builds++;
+		}
 	}
 
-	plist_t filesystem_info_path_node = plist_dict_get_item(filesystem_info_node, "Path");
-	if (!filesystem_info_path_node || plist_get_node_type(filesystem_info_path_node) != PLIST_STRING) {
-		error("ERROR: Unable to find filesystem info path node\n");
-		plist_free(tss_request);
-		return -1;
+	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;
+		}
 	}
-	plist_get_string_val(filesystem_info_path_node, &filesystem);
-	plist_free(tss_request);
 
-	info("Extracting filesystem from IPSW\n");
-	if (ipsw_extract_to_file(ipsw, filesystem, filesystem) < 0) {
-		error("ERROR: Unable to extract filesystem\n");
+	// Extract filesystem from IPSW and return its name
+	char* filesystem = NULL;
+	if (extract_filesystem(client, client->ipsw, build_identity, &filesystem) < 0) {
+		error("ERROR: Unable to extract filesystem from IPSW\n");
+		if (client->tss)
+			plist_free(client->tss);
+		plist_free(buildmanifest);
 		return -1;
 	}
 
-	/* place device into recovery mode if required */
-	if (idevicerestore_mode == NORMAL_MODE) {
+	// if the device is in normal mode, place device into recovery mode
+	if (client->mode->index == MODE_NORMAL) {
 		info("Entering recovery mode...\n");
-		device_error = idevice_new(&device, uuid);
-		if (device_error != IDEVICE_E_SUCCESS) {
-			error("ERROR: Unable to find device\n");
-			plist_free(tss_response);
+		if (normal_enter_recovery(client) < 0) {
+			error("ERROR: Unable to place device into recovery mode\n");
+			if (client->tss)
+				plist_free(client->tss);
+			plist_free(buildmanifest);
 			return -1;
 		}
+	}
 
-		lockdown_error = lockdownd_client_new_with_handshake(device, &lockdown, "idevicerestore");
-		if (lockdown_error != LOCKDOWN_E_SUCCESS) {
-			error("ERROR: Unable to connect to lockdownd service\n");
-			plist_free(tss_response);
-			idevice_free(device);
+	// if the device is in DFU mode, place device into recovery mode
+	if (client->mode->index == MODE_DFU) {
+		if (dfu_enter_recovery(client, build_identity) < 0) {
+			error("ERROR: Unable to place device into recovery mode\n");
+			plist_free(buildmanifest);
+			if (client->tss)
+				plist_free(client->tss);
 			return -1;
 		}
+	}
 
-		lockdown_error = lockdownd_enter_recovery(lockdown);
-		if (lockdown_error != LOCKDOWN_E_SUCCESS) {
-			error("ERROR: Unable to place device in recovery mode\n");
-			lockdownd_client_free(lockdown);
-			plist_free(tss_response);
-			idevice_free(device);
+	// if the device is in recovery mode, place device into restore mode
+	if (client->mode->index == MODE_RECOVERY) {
+		if (recovery_enter_restore(client, build_identity) < 0) {
+			error("ERROR: Unable to place device into restore mode\n");
+			plist_free(buildmanifest);
+			if (client->tss)
+				plist_free(client->tss);
 			return -1;
 		}
+	}
 
-		lockdownd_client_free(lockdown);
-		idevice_free(device);
-		lockdown = NULL;
-		device = NULL;
+	// device is finally in restore mode, let's do this
+	if (client->mode->index == MODE_RESTORE) {
+		info("Restoring device... \n");
+		if (restore_device(client, build_identity, filesystem) < 0) {
+			error("ERROR: Unable to restore device\n");
+			return -1;
+		}
 	}
 
-	/* upload data to make device boot restore mode */
-	if (recovery_send_ibec(ipsw, tss_response) < 0) {
-		error("ERROR: Unable to send iBEC\n");
-		plist_free(tss_response);
-		return -1;
+	info("Cleaning up...\n");
+	if (filesystem)
+		unlink(filesystem);
+
+	info("DONE\n");
+	return 0;
+}
+
+int check_mode(struct idevicerestore_client_t* client) {
+	int mode = MODE_UNKNOWN;
+
+	if (recovery_check_mode() == 0) {
+		mode = MODE_RECOVERY;
 	}
-	sleep(1);
 
-	if (recovery_send_applelogo(ipsw, tss_response) < 0) {
-		error("ERROR: Unable to send AppleLogo\n");
-		plist_free(tss_response);
-		return -1;
+	else if (dfu_check_mode() == 0) {
+		mode = MODE_DFU;
 	}
 
-	if (recovery_send_devicetree(ipsw, tss_response) < 0) {
-		error("ERROR: Unable to send DeviceTree\n");
-		plist_free(tss_response);
-		return -1;
+	else if (normal_check_mode(client->uuid) == 0) {
+		mode = MODE_NORMAL;
 	}
 
-	if (recovery_send_ramdisk(ipsw, tss_response) < 0) {
-		error("ERROR: Unable to send Ramdisk\n");
-		plist_free(tss_response);
-		return -1;
+	else if (restore_check_mode(client->uuid) == 0) {
+		mode = MODE_RESTORE;
 	}
 
-	// for some reason iboot requires a hard reset after ramdisk
-	//   or things start getting wacky
-	printf("Please unplug your device, then plug it back in\n");
-	printf("Hit any key to continue...");
-	getchar();
+	client->mode = &idevicerestore_modes[mode];
+	return mode;
+}
 
-	if (recovery_send_kernelcache(ipsw, tss_response) < 0) {
-		error("ERROR: Unable to send KernelCache\n");
-		plist_free(tss_response);
-		return -1;
-	}
+int check_device(struct idevicerestore_client_t* client) {
+	int device = DEVICE_UNKNOWN;
+	uint32_t bdid = 0;
+	uint32_t cpid = 0;
 
-	idevice_event_subscribe(&device_callback, NULL);
-	info("Waiting for device to enter restore mode\n");
-	// block program until device has entered restore mode
-	while (idevicerestore_mode != RESTORE_MODE) {
-		sleep(1);
-	}
+	switch (client->mode->index) {
+	case MODE_RESTORE:
+		device = restore_check_device(client->uuid);
+		if (device < 0) {
+			device = DEVICE_UNKNOWN;
+		}
+		break;
 
-	device_error = idevice_new(&device, uuid);
-	if (device_error != IDEVICE_E_SUCCESS) {
-		error("ERROR: Unable to open device\n");
-		plist_free(tss_response);
-		return -1;
-	}
+	case MODE_NORMAL:
+		device = normal_check_device(client->uuid);
+		if (device < 0) {
+			device = DEVICE_UNKNOWN;
+		}
+		break;
 
-	restored_client_t restore = NULL;
-	restored_error_t restore_error = restored_client_new(device, &restore, "idevicerestore");
-	if (restore_error != RESTORE_E_SUCCESS) {
-		error("ERROR: Unable to start restored client\n");
-		plist_free(tss_response);
-		idevice_free(device);
-		return -1;
-	}
+	case MODE_DFU:
+	case MODE_RECOVERY:
+		if (get_cpid(client, &cpid) < 0) {
+			error("ERROR: Unable to get device CPID\n");
+			break;
+		}
 
-	char* type = NULL;
-	uint64_t version = 0;
-	if (restored_query_type(restore, &type, &version) != RESTORE_E_SUCCESS) {
-		error("ERROR: Device is not in restore mode. QueryType returned \"%s\"\n", type);
-		plist_free(tss_response);
-		restored_client_free(restore);
-		idevice_free(device);
-		return -1;
-	}
-	info("Device has successfully entered restore mode\n");
-
-	/* start restore process */
-	char* kernelcache = NULL;
-	info("Restore protocol version is %llu.\n", version);
-	restore_error = restored_start_restore(restore);
-	if (restore_error == RESTORE_E_SUCCESS) {
-		while (!idevicerestore_quit) {
-			plist_t message = NULL;
-			restore_error = restored_receive(restore, &message);
-			plist_t msgtype_node = plist_dict_get_item(message, "MsgType");
-			if (msgtype_node && PLIST_STRING == plist_get_node_type(msgtype_node)) {
-				char *msgtype = NULL;
-				plist_get_string_val(msgtype_node, &msgtype);
-				if (!strcmp(msgtype, "ProgressMsg")) {
-					restore_error = restore_handle_progress_msg(restore, message);
-
-				} else if (!strcmp(msgtype, "DataRequestMsg")) {
-					// device is requesting data to be sent
-					plist_t datatype_node = plist_dict_get_item(message, "DataType");
-					if (datatype_node && PLIST_STRING == plist_get_node_type(datatype_node)) {
-						char *datatype = NULL;
-						plist_get_string_val(datatype_node, &datatype);
-						if (!strcmp(datatype, "SystemImageData")) {
-							asr_send_system_image_data_from_file(device, restore, filesystem);
-
-						} else if (!strcmp(datatype, "KernelCache")) {
-							int kernelcache_size = 0;
-							char* kernelcache_data = NULL;
-							if (get_signed_component_by_name(ipsw, tss_response, "KernelCache", &kernelcache_data, &kernelcache_size) < 0) {
-								error("ERROR: Unable to get kernelcache file\n");
-								return -1;
-							}
-							restore_send_kernelcache(restore, kernelcache_data, kernelcache_size);
-							free(kernelcache_data);
-
-						} else if (!strcmp(datatype, "NORData")) {
-							restore_send_nor_data(restore, ipsw, tss_response);
-
-						} else {
-							// Unknown DataType!!
-							error("Unknown DataType\n");
-							return -1;
-						}
-					}
-
-				} else if (!strcmp(msgtype, "StatusMsg")) {
-					restore_error = restore_handle_status_msg(restore, message);
-
-				} else {
-					info("Received unknown message type: %s\n", msgtype);
-				}
+		switch (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(client, &bdid) < 0) {
+				error("ERROR: Unable to get device BDID\n");
+				break;
 			}
 
-			if (RESTORE_E_SUCCESS != restore_error) {
-				error("Invalid return status %d\n", restore_error);
-				//idevicerestore_quit = 1;
+			switch (bdid) {
+			case BDID_IPHONE2G:
+				device = DEVICE_IPHONE2G;
+				break;
+
+			case BDID_IPHONE3G:
+				device = DEVICE_IPHONE3G;
+				break;
+
+			case BDID_IPOD1G:
+				device = DEVICE_IPOD1G;
+				break;
+
+			default:
+				device = DEVICE_UNKNOWN;
+				break;
 			}
+			break;
+
+		case CPID_IPHONE3GS:
+			device = DEVICE_IPHONE3GS;
+			break;
+
+		case CPID_IPOD2G:
+			device = DEVICE_IPOD2G;
+			break;
+
+		case CPID_IPOD3G:
+			device = DEVICE_IPOD3G;
+			break;
 
-			plist_free(message);
+		case CPID_IPAD1G:
+			device = DEVICE_IPAD1G;
+			break;
+
+		default:
+			device = DEVICE_UNKNOWN;
+			break;
 		}
-	} else {
-		error("ERROR: Could not start restore. %d\n", restore_error);
-	}
+		break;
 
-	restored_client_free(restore);
-	plist_free(tss_response);
-	idevice_free(device);
-	unlink(filesystem);
-	return 0;
-}
+	default:
+		device = DEVICE_UNKNOWN;
+		break;
 
-void device_callback(const idevice_event_t* event, void *user_data) {
-	if (event->event == IDEVICE_DEVICE_ADD) {
-		idevicerestore_mode = RESTORE_MODE;
-	} else if(event->event == IDEVICE_DEVICE_REMOVE) {
-		idevicerestore_quit = 1;
 	}
-}
 
-void usage(int argc, char* argv[]) {
-	char *name = strrchr(argv[0], '/');
-	printf("Usage: %s [OPTIONS] FILE\n", (name ? name + 1 : argv[0]));
-	printf("Restore/upgrade IPSW firmware FILE to an iPhone/iPod Touch.\n");
-	printf("  -d, \t\tenable communication debugging\n");
-	printf("  -u, \t\ttarget specific device by its 40-digit device UUID\n");
-	printf("  -h, \t\tprints usage information\n");
-	printf("  -c, \t\trestore with a custom firmware\n");
-	printf("  -v, \t\tenable incremental levels of verboseness\n");
-	printf("\n");
-	exit(1);
+	client->device = &idevicerestore_devices[device];
+	return device;
 }
 
-int write_file(const char* filename, char* data, int size) {
-	debug("Writing data to %s\n", filename);
-	FILE* file = fopen(filename, "wb");
-	if (file == NULL) {
-		error("read_file: Unable to open file %s\n", filename);
-		return -1;
-	}
+int get_bdid(struct idevicerestore_client_t* client, uint32_t* bdid) {
+	switch (client->mode->index) {
+	case MODE_NORMAL:
+		if (normal_get_bdid(client->uuid, bdid) < 0) {
+			*bdid = 0;
+			return -1;
+		}
+		break;
 
-	int bytes = fwrite(data, 1, size, file);
-	fclose(file);
+	case MODE_DFU:
+	case MODE_RECOVERY:
+		if (recovery_get_bdid(client, bdid) < 0) {
+			*bdid = 0;
+			return -1;
+		}
+		break;
 
-	if (bytes != size) {
-		error("ERROR: Unable to write entire file: %s: %d %d\n", filename, bytes, size);
+	default:
+		error("ERROR: Device is in an invalid state\n");
 		return -1;
 	}
 
-	return size;
+	return 0;
 }
 
-int get_tss_data_by_path(plist_t tss, const char* path, char** pname, char** pblob) {
-	*pname = NULL;
-	*pblob = NULL;
-
-	char* key = NULL;
-	plist_t tss_entry = NULL;
-	plist_dict_iter iter = NULL;
-	plist_dict_new_iter(tss, &iter);
-	while (1) {
-		plist_dict_next_item(tss, iter, &key, &tss_entry);
-		if (key == NULL)
-			break;
-
-		if (!tss_entry || plist_get_node_type(tss_entry) != PLIST_DICT) {
-			continue;
+int get_cpid(struct idevicerestore_client_t* client, uint32_t* cpid) {
+	switch (client->mode->index) {
+	case MODE_NORMAL:
+		if (normal_get_cpid(client->uuid, cpid) < 0) {
+			client->device->chip_id = -1;
+			return -1;
 		}
+		break;
 
-		char* entry_path = NULL;
-		plist_t entry_path_node = plist_dict_get_item(tss_entry, "Path");
-		if (!entry_path_node || plist_get_node_type(entry_path_node) != PLIST_STRING) {
-			error("ERROR: Unable to find TSS path node in entry %s\n", key);
+	case MODE_DFU:
+	case MODE_RECOVERY:
+		if (recovery_get_cpid(client, cpid) < 0) {
+			client->device->chip_id = -1;
 			return -1;
 		}
-		plist_get_string_val(entry_path_node, &entry_path);
-		if (strcmp(path, entry_path) == 0) {
-			char* blob = NULL;
-			uint64_t blob_size = 0;
-			plist_t blob_node = plist_dict_get_item(tss_entry, "Blob");
-			if (!blob_node || plist_get_node_type(blob_node) != PLIST_DATA) {
-				error("ERROR: Unable to find TSS blob node in entry %s\n", key);
-				return -1;
-			}
-			plist_get_data_val(blob_node, &blob, &blob_size);
-			*pname = key;
-			*pblob = blob;
-			return 0;
-		}
+		break;
 
-		free(key);
+	default:
+		error("ERROR: Device is in an invalid state\n");
+		return -1;
 	}
-	plist_free(tss_entry);
 
-	return -1;
+	return 0;
 }
 
-int get_tss_data_by_name(plist_t tss, const char* entry, char** ppath, char** pblob) {
-	*ppath = NULL;
-	*pblob = NULL;
-
-	plist_t node = plist_dict_get_item(tss, entry);
-	if (!node || plist_get_node_type(node) != PLIST_DICT) {
-		error("ERROR: Unable to find %s entry in TSS response\n", entry);
-		return -1;
+int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
+	if(client->device->index <= DEVICE_IPOD2G) {
+		*ecid = 0;
+		return 0;
 	}
 
-	char* path = NULL;
-	plist_t path_node = plist_dict_get_item(node, "Path");
-	if (!path_node || plist_get_node_type(path_node) != PLIST_STRING) {
-		error("ERROR: Unable to find %s path in entry\n", path);
-		return -1;
-	}
-	plist_get_string_val(path_node, &path);
+	switch (client->mode->index) {
+	case MODE_NORMAL:
+		if (normal_get_ecid(client->uuid, ecid) < 0) {
+			*ecid = 0;
+			return -1;
+		}
+		break;
 
-	char* blob = NULL;
-	uint64_t blob_size = 0;
-	plist_t blob_node = plist_dict_get_item(node, "Blob");
-	if (!blob_node || plist_get_node_type(blob_node) != PLIST_DATA) {
-		error("ERROR: Unable to find %s blob in entry\n", path);
-		free(path);
+	case MODE_DFU:
+	case MODE_RECOVERY:
+		if (recovery_get_ecid(client, ecid) < 0) {
+			*ecid = 0;
+			return -1;
+		}
+		break;
+
+	default:
+		error("ERROR: Device is in an invalid state\n");
 		return -1;
 	}
-	plist_get_data_val(blob_node, &blob, &blob_size);
 
-	*ppath = path;
-	*pblob = blob;
 	return 0;
 }
 
-int get_signed_component_by_name(char* ipsw, plist_t tss, char* component, char** pdata, int* psize) {
+int extract_buildmanifest(struct idevicerestore_client_t* client, const char* ipsw, plist_t* buildmanifest) {
 	int size = 0;
 	char* data = NULL;
-	char* path = NULL;
-	char* blob = NULL;
-	img3_file* img3 = NULL;
-	irecv_error_t error = 0;
+	int device = client->device->index;
 
-	info("Extracting %s from TSS response\n", component);
-	if (get_tss_data_by_name(tss, component, &path, &blob) < 0) {
-		error("ERROR: Unable to get data for TSS %s entry\n", component);
-		return -1;
+	/* older devices don't require personalized firmwares and use a BuildManifesto.plist */
+	if (ipsw_extract_to_memory(ipsw, "BuildManifesto.plist", &data, &size) == 0) {
+		plist_from_xml(data, size, buildmanifest);
+		return 0;
 	}
 
-	info("Extracting %s from %s\n", path, ipsw);
-	if (ipsw_extract_to_memory(ipsw, path, &data, &size) < 0) {
-		error("ERROR: Unable to extract %s from %s\n", path, ipsw);
-		free(path);
-		free(blob);
-		return -1;
+	data = NULL;
+	size = 0;
+
+	/* whereas newer devices do not require personalized firmwares and use a BuildManifest.plist */
+	if (ipsw_extract_to_memory(ipsw, "BuildManifest.plist", &data, &size) == 0) {
+		plist_from_xml(data, size, buildmanifest);
+		return 0;
 	}
 
-	img3 = img3_parse_file(data, size);
-	if (img3 == NULL) {
-		error("ERROR: Unable to parse IMG3: %s\n", path);
-		free(data);
-		free(path);
-		free(blob);
-		return -1;
+	return -1;
+}
+
+plist_t get_build_identity(struct idevicerestore_client_t* client, plist_t buildmanifest, uint32_t identity) {
+	// fetch build identities array from BuildManifest
+	plist_t build_identities_array = plist_dict_get_item(buildmanifest, "BuildIdentities");
+	if (!build_identities_array || plist_get_node_type(build_identities_array) != PLIST_ARRAY) {
+		error("ERROR: Unable to find build identities node\n");
+		return NULL;
 	}
-	if (data) {
-		free(data);
-		data = NULL;
+
+	// check and make sure this identity exists in buildmanifest
+	if (identity >= plist_array_get_size(build_identities_array)) {
+		return NULL;
 	}
 
-	if (idevicerestore_custom == 0) {
-		if (img3_replace_signature(img3, blob) < 0) {
-			error("ERROR: Unable to replace IMG3 signature\n");
-			free(path);
-			free(blob);
-			return -1;
-		}
+	plist_t build_identity = plist_array_get_item(build_identities_array, identity);
+	if (!build_identity || plist_get_node_type(build_identity) != PLIST_DICT) {
+		error("ERROR: Unable to find build identities node\n");
+		return NULL;
 	}
 
-	if (img3_get_data(img3, &data, &size) < 0) {
-		error("ERROR: Unable to reconstruct IMG3\n");
-		img3_free(img3);
-		free(path);
+	return plist_copy(build_identity);
+}
+
+int get_shsh_blobs(struct idevicerestore_client_t* client, uint64_t ecid, plist_t build_identity, plist_t* tss) {
+	plist_t request = NULL;
+	plist_t response = NULL;
+	*tss = NULL;
+
+	request = tss_create_request(build_identity, ecid);
+	if (request == NULL) {
+		error("ERROR: Unable to create TSS request\n");
 		return -1;
 	}
 
-	if (idevicerestore_debug) {
-		char* out = strrchr(path, '/');
-		if (out != NULL) {
-			out++;
-		} else {
-			out = path;
-		}
-		write_file(out, data, size);
+	info("Sending TSS request\n");
+	response = tss_send_request(request);
+	if (response == NULL) {
+		plist_free(request);
+		return -1;
 	}
 
-	if (img3) {
-		img3_free(img3);
-		img3 = NULL;
-	}
-	if (blob) {
-		free(blob);
-		blob = NULL;
-	}
-	if (path) {
-		free(path);
-		path = NULL;
+	plist_free(request);
+	*tss = response;
+	return 0;
+}
+
+int get_build_count(plist_t buildmanifest) {
+	// fetch build identities array from BuildManifest
+	plist_t build_identities_array = plist_dict_get_item(buildmanifest, "BuildIdentities");
+	if (!build_identities_array || plist_get_node_type(build_identities_array) != PLIST_ARRAY) {
+		error("ERROR: Unable to find build identities node\n");
+		return -1;
 	}
 
-	*pdata = data;
-	*psize = size;
-	return 0;
+	// check and make sure this identity exists in buildmanifest
+	return plist_array_get_size(build_identities_array);
 }
 
-int get_signed_component_by_path(char* ipsw, plist_t tss, char* path, char** pdata, int* psize) {
-	int size = 0;
-	char* data = NULL;
-	char* name = NULL;
-	char* blob = NULL;
-	img3_file* img3 = NULL;
-	irecv_error_t error = 0;
+int extract_filesystem(struct idevicerestore_client_t* client, const char* ipsw, plist_t build_identity, char** filesystem) {
+	char* filename = NULL;
 
-	info("Extracting %s from TSS response\n", path);
-	if (get_tss_data_by_path(tss, path, &name, &blob) < 0) {
-		error("ERROR: Unable to get data for TSS %s entry\n", path);
+	if (build_identity_get_component_path(build_identity, "OS", &filename) < 0) {
+		error("ERROR: Unable get path for filesystem component\n");
 		return -1;
 	}
 
-	info("Extracting %s from %s\n", path, ipsw);
-	if (ipsw_extract_to_memory(ipsw, path, &data, &size) < 0) {
-		error("ERROR: Unable to extract %s from %s\n", path, ipsw);
-		free(path);
-		free(blob);
+	info("Extracting filesystem from IPSW\n");
+	if (ipsw_extract_to_file(ipsw, filename, filename) < 0) {
+		error("ERROR: Unable to extract filesystem\n");
 		return -1;
 	}
 
-	img3 = img3_parse_file(data, size);
-	if (img3 == NULL) {
-		error("ERROR: Unable to parse IMG3: %s\n", path);
-		free(data);
-		free(path);
-		free(blob);
+	*filesystem = filename;
+	return 0;
+}
+
+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;
+	char* component_blob = NULL;
+	char* component_name = NULL;
+
+	component_name = strrchr(path, '/');
+	if (component_name != NULL)
+		component_name++;
+	else
+		component_name = (char*) path;
+
+	info("Extracting %s\n", component_name);
+	if (ipsw_extract_to_memory(ipsw, path, &component_data, &component_size) < 0) {
+		error("ERROR: Unable to extract %s from %s\n", component_name, ipsw);
 		return -1;
 	}
-	if (data) {
-		free(data);
-		data = NULL;
-	}
 
-	if (idevicerestore_custom == 0) {
-		if (img3_replace_signature(img3, blob) < 0) {
+	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);
+
+		/* 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 (img3_replace_signature(img3, component_blob) < 0) {
 			error("ERROR: Unable to replace IMG3 signature\n");
-			free(name);
-			free(blob);
+			free(component_blob);
+			img3_free(img3);
 			return -1;
 		}
-	}
 
-	if (img3_get_data(img3, &data, &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);
-		free(name);
-		return -1;
 	}
 
 	if (idevicerestore_debug) {
-		char* out = strrchr(path, '/');
-		if (out != NULL) {
-			out++;
-		} else {
-			out = path;
-		}
-		write_file(out, data, size);
+		write_file(component_name, component_data, component_size);
 	}
 
-	if (img3) {
-		img3_free(img3);
-		img3 = NULL;
+	*data = component_data;
+	*size = component_size;
+	return 0;
+}
+
+int build_identity_get_component_path(plist_t build_identity, const char* component, char** path) {
+	char* filename = NULL;
+
+	plist_t manifest_node = plist_dict_get_item(build_identity, "Manifest");
+	if (!manifest_node || plist_get_node_type(manifest_node) != PLIST_DICT) {
+		error("ERROR: Unable to find manifest node\n");
+		if (filename)
+			free(filename);
+		return -1;
+	}
+
+	plist_t component_node = plist_dict_get_item(manifest_node, component);
+	if (!component_node || plist_get_node_type(component_node) != PLIST_DICT) {
+		error("ERROR: Unable to find component node for %s\n", component);
+		if (filename)
+			free(filename);
+		return -1;
 	}
-	if (blob) {
-		free(blob);
-		blob = NULL;
+
+	plist_t component_info_node = plist_dict_get_item(component_node, "Info");
+	if (!component_info_node || plist_get_node_type(component_info_node) != PLIST_DICT) {
+		error("ERROR: Unable to find component info node for %s\n", component);
+		if (filename)
+			free(filename);
+		return -1;
 	}
-	if (path) {
-		free(name);
-		name = NULL;
+
+	plist_t component_info_path_node = plist_dict_get_item(component_info_node, "Path");
+	if (!component_info_path_node || plist_get_node_type(component_info_path_node) != PLIST_STRING) {
+		error("ERROR: Unable to find component info path node for %s\n", component);
+		if (filename)
+			free(filename);
+		return -1;
 	}
+	plist_get_string_val(component_info_path_node, &filename);
 
-	*pdata = data;
-	*psize = size;
+	*path = filename;
 	return 0;
 }
+
diff --git a/src/idevicerestore.h b/src/idevicerestore.h
index 3dcf1d5..f529b5b 100644
--- a/src/idevicerestore.h
+++ b/src/idevicerestore.h
@@ -1,5 +1,5 @@
 /*
- * idevicerestore.g
+ * idevicerestore.h
  * Restore device firmware and filesystem
  *
  * Copyright (c) 2010 Joshua Hill. All Rights Reserved.
@@ -22,10 +22,31 @@
 #ifndef IDEVICERESTORE_H
 #define IDEVICERESTORE_H
 
-#define error(...) fprintf(stderr, __VA_ARGS__)
-#define info(...) if(idevicerestore_debug >= 1) fprintf(stderr, __VA_ARGS__)
-#define debug(...) if(idevicerestore_debug >= 2) fprintf(stderr, __VA_ARGS__)
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stdint.h>
+#include <plist/plist.h>
+
+#include "common.h"
 
-extern int idevicerestore_debug;
+void usage(int argc, char* argv[]);
+int check_mode(struct idevicerestore_client_t* client);
+int check_device(struct idevicerestore_client_t* client);
+int get_build_count(plist_t buildmanifest);
+int get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid);
+int get_bdid(struct idevicerestore_client_t* client, uint32_t* bdid);
+int get_cpid(struct idevicerestore_client_t* client, uint32_t* cpid);
+int extract_buildmanifest(struct idevicerestore_client_t* client, const char* ipsw, plist_t* buildmanifest);
+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 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
+}
+#endif
 
 #endif
diff --git a/src/img3.c b/src/img3.c
index ec7bcbf..1969ece 100644
--- a/src/img3.c
+++ b/src/img3.c
@@ -24,25 +24,26 @@
 #include <string.h>
 
 #include "img3.h"
+#include "common.h"
 #include "idevicerestore.h"
 
 img3_file* img3_parse_file(char* data, int size) {
 	int data_offset = 0;
 	img3_header* header = (img3_header*) data;
-	if(header->signature != kImg3Container) {
+	if (header->signature != kImg3Container) {
 		error("ERROR: Invalid IMG3 file\n");
 		return NULL;
 	}
 
 	img3_file* image = (img3_file*) malloc(sizeof(img3_file));
-	if(image == NULL) {
+	if (image == NULL) {
 		error("ERROR: Unable to allocate memory for IMG3 file\n");
 		return NULL;
 	}
 	memset(image, '\0', sizeof(img3_file));
 
 	image->header = (img3_header*) malloc(sizeof(img3_header));
-	if(image->header == NULL) {
+	if (image->header == NULL) {
 		error("ERROR: Unable to allocate memory for IMG3 header\n");
 		img3_free(image);
 		return NULL;
@@ -51,12 +52,12 @@ img3_file* img3_parse_file(char* data, int size) {
 	data_offset += sizeof(img3_header);
 
 	img3_element_header* current = NULL;
-	while(data_offset < size) {
+	while (data_offset < size) {
 		current = (img3_element_header*) &data[data_offset];
-		switch(current->signature) {
+		switch (current->signature) {
 		case kTypeElement:
 			image->type_element = img3_parse_element(&data[data_offset]);
-			if(image->type_element == NULL) {
+			if (image->type_element == NULL) {
 				error("ERROR: Unable to parse TYPE element\n");
 				img3_free(image);
 				return NULL;
@@ -66,7 +67,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kDataElement:
 			image->data_element = img3_parse_element(&data[data_offset]);
-			if(image->data_element == NULL) {
+			if (image->data_element == NULL) {
 				error("ERROR: Unable to parse DATA element\n");
 				img3_free(image);
 				return NULL;
@@ -76,7 +77,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kVersElement:
 			image->vers_element = img3_parse_element(&data[data_offset]);
-			if(image->vers_element == NULL) {
+			if (image->vers_element == NULL) {
 				error("ERROR: Unable to parse VERS element\n");
 				img3_free(image);
 				return NULL;
@@ -86,7 +87,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kSepoElement:
 			image->sepo_element = img3_parse_element(&data[data_offset]);
-			if(image->sepo_element == NULL) {
+			if (image->sepo_element == NULL) {
 				error("ERROR: Unable to parse SEPO element\n");
 				img3_free(image);
 				return NULL;
@@ -96,7 +97,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kBordElement:
 			image->bord_element = img3_parse_element(&data[data_offset]);
-			if(image->bord_element == NULL) {
+			if (image->bord_element == NULL) {
 				error("ERROR: Unable to parse BORD element\n");
 				img3_free(image);
 				return NULL;
@@ -105,9 +106,9 @@ img3_file* img3_parse_file(char* data, int size) {
 			break;
 
 		case kKbagElement:
-			if(image->kbag1_element == NULL) {
+			if (image->kbag1_element == NULL) {
 				image->kbag1_element = img3_parse_element(&data[data_offset]);
-				if(image->kbag1_element == NULL) {
+				if (image->kbag1_element == NULL) {
 					error("ERROR: Unable to parse first KBAG element\n");
 					img3_free(image);
 					return NULL;
@@ -115,7 +116,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 			} else {
 				image->kbag2_element = img3_parse_element(&data[data_offset]);
-				if(image->kbag2_element == NULL) {
+				if (image->kbag2_element == NULL) {
 					error("ERROR: Unable to parse second KBAG element\n");
 					img3_free(image);
 					return NULL;
@@ -126,7 +127,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kEcidElement:
 			image->ecid_element = img3_parse_element(&data[data_offset]);
-			if(image->ecid_element == NULL) {
+			if (image->ecid_element == NULL) {
 				error("ERROR: Unable to parse ECID element\n");
 				img3_free(image);
 				return NULL;
@@ -136,7 +137,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kShshElement:
 			image->shsh_element = img3_parse_element(&data[data_offset]);
-			if(image->shsh_element == NULL) {
+			if (image->shsh_element == NULL) {
 				error("ERROR: Unable to parse SHSH element\n");
 				img3_free(image);
 				return NULL;
@@ -146,7 +147,7 @@ img3_file* img3_parse_file(char* data, int size) {
 
 		case kCertElement:
 			image->cert_element = img3_parse_element(&data[data_offset]);
-			if(image->cert_element == NULL) {
+			if (image->cert_element == NULL) {
 				error("ERROR: Unable to parse CERT element\n");
 				img3_free(image);
 				return NULL;
@@ -154,6 +155,16 @@ img3_file* img3_parse_file(char* data, int size) {
 			debug("Parsed CERT element\n");
 			break;
 
+		case kUnknElement:
+			image->unkn_element = img3_parse_element(&data[data_offset]);
+			if (image->unkn_element == NULL) {
+				error("ERROR: Unable to parse UNKN element\n");
+				img3_free(image);
+				return NULL;
+			}
+			debug("Parsed UNKN element\n");
+			break;
+
 		default:
 			error("ERROR: Unknown IMG3 element type\n");
 			img3_free(image);
@@ -168,14 +179,14 @@ img3_file* img3_parse_file(char* data, int size) {
 img3_element* img3_parse_element(char* data) {
 	img3_element_header* element_header = (img3_element_header*) data;
 	img3_element* element = (img3_element*) malloc(sizeof(img3_element));
-	if(element == NULL) {
+	if (element == NULL) {
 		error("ERROR: Unable to allocate memory for IMG3 element\n");
 		return NULL;
 	}
 	memset(element, '\0', sizeof(img3_element));
 
 	element->data = (char*) malloc(element_header->full_size);
-	if(element->data == NULL) {
+	if (element->data == NULL) {
 		error("ERROR: Unable to allocate memory for IMG3 element data\n");
 		free(element);
 		return NULL;
@@ -188,69 +199,74 @@ img3_element* img3_parse_element(char* data) {
 }
 
 void img3_free(img3_file* image) {
-	if(image != NULL) {
-		if(image->header != NULL) {
+	if (image != NULL) {
+		if (image->header != NULL) {
 			free(image->header);
 		}
 
-		if(image->type_element != NULL) {
+		if (image->type_element != NULL) {
 			img3_free_element(image->type_element);
 			image->type_element = NULL;
 		}
 
-		if(image->data_element != NULL) {
+		if (image->data_element != NULL) {
 			img3_free_element(image->data_element);
 			image->data_element = NULL;
 		}
 
-		if(image->vers_element != NULL) {
+		if (image->vers_element != NULL) {
 			img3_free_element(image->vers_element);
 			image->vers_element = NULL;
 		}
 
-		if(image->sepo_element != NULL) {
+		if (image->sepo_element != NULL) {
 			img3_free_element(image->sepo_element);
 			image->sepo_element = NULL;
 		}
 
-		if(image->bord_element != NULL) {
+		if (image->bord_element != NULL) {
 			img3_free_element(image->bord_element);
 			image->bord_element = NULL;
 		}
 
-		if(image->kbag1_element != NULL) {
+		if (image->kbag1_element != NULL) {
 			img3_free_element(image->kbag1_element);
 			image->kbag1_element = NULL;
 		}
 
-		if(image->kbag2_element != NULL) {
+		if (image->kbag2_element != NULL) {
 			img3_free_element(image->kbag2_element);
 			image->kbag2_element = NULL;
 		}
 
-		if(image->ecid_element != NULL) {
+		if (image->ecid_element != NULL) {
 			img3_free_element(image->ecid_element);
 			image->ecid_element = NULL;
 		}
 
-		if(image->shsh_element != NULL) {
+		if (image->shsh_element != NULL) {
 			img3_free_element(image->shsh_element);
 			image->shsh_element = NULL;
 		}
 
-		if(image->cert_element != NULL) {
+		if (image->cert_element != NULL) {
 			img3_free_element(image->cert_element);
 			image->cert_element = NULL;
 		}
 
+		if (image->unkn_element != NULL) {
+			img3_free_element(image->unkn_element);
+			image->unkn_element = NULL;
+		}
+
 		free(image);
 		image = NULL;
 	}
 }
 
 void img3_free_element(img3_element* element) {
-	if(element != NULL) {
-		if(element->data != NULL) {
+	if (element != NULL) {
+		if (element->data != NULL) {
 			free(element->data);
 			element->data = NULL;
 		}
@@ -262,37 +278,37 @@ void img3_free_element(img3_element* element) {
 int img3_replace_signature(img3_file* image, char* signature) {
 	int offset = 0;
 	img3_element* ecid = img3_parse_element(&signature[offset]);
-	if(ecid == NULL || ecid->type != kEcidElement) {
+	if (ecid == NULL || ecid->type != kEcidElement) {
 		error("ERROR: Unable to find ECID element in signature\n");
 		return -1;
 	}
 	offset += ecid->header->full_size;
 
 	img3_element* shsh = img3_parse_element(&signature[offset]);
-	if(shsh == NULL || shsh->type != kShshElement) {
+	if (shsh == NULL || shsh->type != kShshElement) {
 		error("ERROR: Unable to find SHSH element in signature\n");
 		return -1;
 	}
 	offset += shsh->header->full_size;
 
 	img3_element* cert = img3_parse_element(&signature[offset]);
-	if(cert == NULL || cert->type != kCertElement) {
+	if (cert == NULL || cert->type != kCertElement) {
 		error("ERROR: Unable to find CERT element in signature\n");
 		return -1;
 	}
 	offset += cert->header->full_size;
 
-	if(image->ecid_element != NULL) {
+	if (image->ecid_element != NULL) {
 		img3_free_element(image->ecid_element);
 	}
 	image->ecid_element = ecid;
 
-	if(image->shsh_element != NULL) {
+	if (image->shsh_element != NULL) {
 		img3_free_element(image->shsh_element);
 	}
 	image->shsh_element = shsh;
 
-	if(image->cert_element != NULL) {
+	if (image->cert_element != NULL) {
 		img3_free_element(image->cert_element);
 	}
 	image->cert_element = cert;
@@ -305,39 +321,42 @@ int img3_get_data(img3_file* image, char** pdata, int* psize) {
 	int size = sizeof(img3_header);
 
 	// Add up the size of the image first so we can allocate our memory
-	if(image->type_element != NULL) {
+	if (image->type_element != NULL) {
 		size += image->type_element->header->full_size;
 	}
-	if(image->data_element != NULL) {
+	if (image->data_element != NULL) {
 		size += image->data_element->header->full_size;
 	}
-	if(image->vers_element != NULL) {
+	if (image->vers_element != NULL) {
 		size += image->vers_element->header->full_size;
 	}
-	if(image->sepo_element != NULL) {
+	if (image->sepo_element != NULL) {
 		size += image->sepo_element->header->full_size;
 	}
-	if(image->bord_element != NULL) {
+	if (image->bord_element != NULL) {
 		size += image->bord_element->header->full_size;
 	}
-	if(image->kbag1_element != NULL) {
+	if (image->kbag1_element != NULL) {
 		size += image->kbag1_element->header->full_size;
 	}
-	if(image->kbag2_element != NULL) {
+	if (image->kbag2_element != NULL) {
 		size += image->kbag2_element->header->full_size;
 	}
-	if(image->ecid_element != NULL) {
+	if (image->ecid_element != NULL) {
 		size += image->ecid_element->header->full_size;
 	}
-	if(image->shsh_element != NULL) {
+	if (image->shsh_element != NULL) {
 		size += image->shsh_element->header->full_size;
 	}
-	if(image->cert_element != NULL) {
+	if (image->cert_element != NULL) {
 		size += image->cert_element->header->full_size;
 	}
+	if (image->unkn_element != NULL) {
+		size += image->unkn_element->header->full_size;
+	}
 
 	char* data = (char*) malloc(size);
-	if(data == NULL) {
+	if (data == NULL) {
 		error("ERROR: Unable to allocate memory for IMG3 data\n");
 		return -1;
 	}
@@ -351,49 +370,53 @@ int img3_get_data(img3_file* image, char** pdata, int* psize) {
 	offset += sizeof(img3_header);
 
 	// Copy each section over to the new buffer
-	if(image->type_element != NULL) {
+	if (image->type_element != NULL) {
 		memcpy(&data[offset], image->type_element->data, image->type_element->header->full_size);
 		offset += image->type_element->header->full_size;
 	}
-	if(image->data_element != NULL) {
+	if (image->data_element != NULL) {
 		memcpy(&data[offset], image->data_element->data, image->data_element->header->full_size);
 		offset += image->data_element->header->full_size;
 	}
-	if(image->vers_element != NULL) {
+	if (image->vers_element != NULL) {
 		memcpy(&data[offset], image->vers_element->data, image->vers_element->header->full_size);
 		offset += image->vers_element->header->full_size;
 	}
-	if(image->sepo_element != NULL) {
+	if (image->sepo_element != NULL) {
 		memcpy(&data[offset], image->sepo_element->data, image->sepo_element->header->full_size);
 		offset += image->sepo_element->header->full_size;
 	}
-	if(image->bord_element != NULL) {
+	if (image->bord_element != NULL) {
 		memcpy(&data[offset], image->bord_element->data, image->bord_element->header->full_size);
 		offset += image->bord_element->header->full_size;
 	}
-	if(image->kbag1_element != NULL) {
+	if (image->kbag1_element != NULL) {
 		memcpy(&data[offset], image->kbag1_element->data, image->kbag1_element->header->full_size);
 		offset += image->kbag1_element->header->full_size;
 	}
-	if(image->kbag2_element != NULL) {
+	if (image->kbag2_element != NULL) {
 		memcpy(&data[offset], image->kbag2_element->data, image->kbag2_element->header->full_size);
 		offset += image->kbag2_element->header->full_size;
 	}
-	if(image->ecid_element != NULL) {
+	if (image->ecid_element != NULL) {
 		memcpy(&data[offset], image->ecid_element->data, image->ecid_element->header->full_size);
 		offset += image->ecid_element->header->full_size;
 	}
-	if(image->shsh_element != NULL) {
+	if (image->shsh_element != NULL) {
 		memcpy(&data[offset], image->shsh_element->data, image->shsh_element->header->full_size);
 		header->shsh_offset = offset - sizeof(img3_header);
 		offset += image->shsh_element->header->full_size;
 	}
-	if(image->cert_element != NULL) {
+	if (image->cert_element != NULL) {
 		memcpy(&data[offset], image->cert_element->data, image->cert_element->header->full_size);
 		offset += image->cert_element->header->full_size;
 	}
+	if (image->unkn_element != NULL) {
+		memcpy(&data[offset], image->unkn_element->data, image->unkn_element->header->full_size);
+		offset += image->unkn_element->header->full_size;
+	}
 
-	if(offset != size) {
+	if (offset != size) {
 		error("ERROR: Incorrectly sized image data\n");
 		free(data);
 		*pdata = 0;
diff --git a/src/img3.h b/src/img3.h
index a19ae99..cb042b5 100644
--- a/src/img3.h
+++ b/src/img3.h
@@ -19,43 +19,48 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef IMG3_H
-#define IMG3_H
+#ifndef IDEVICERESTORE_IMG3_H
+#define IDEVICERESTORE_IMG3_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 typedef enum {
-    kNorContainer  = 0x696D6733, // img3
-    kImg3Container = 0x496D6733, // Img3
-    k8900Container = 0x30303938, // 8900
-    kImg2Container = 0x494D4732  // IMG2
+	kNorContainer = 0x696D6733,  // img3
+	kImg3Container = 0x496D6733, // Img3
+	k8900Container = 0x30303938, // 8900
+	kImg2Container = 0x494D4732  // IMG2
 } img3_container;
 
 typedef enum {
-    kDataElement = 0x44415441, // DATA
-    kTypeElement = 0x54595045, // TYPE
-    kKbagElement = 0x4B424147, // KBAG
-    kShshElement = 0x53485348, // SHSH
-    kCertElement = 0x43455254, // CERT
-    kChipElement = 0x43484950, // CHIP
-    kProdElement = 0x50524F44, // PROD
-    kSdomElement = 0x53444F4D, // SDOM
-    kVersElement = 0x56455253, // VERS
-    kBordElement = 0x424F5244, // BORD
-    kSepoElement = 0x5345504F, // SEPO
-    kEcidElement = 0x45434944  // ECID
+	kDataElement = 0x44415441, // DATA
+	kTypeElement = 0x54595045, // TYPE
+	kKbagElement = 0x4B424147, // KBAG
+	kShshElement = 0x53485348, // SHSH
+	kCertElement = 0x43455254, // CERT
+	kChipElement = 0x43484950, // CHIP
+	kProdElement = 0x50524F44, // PROD
+	kSdomElement = 0x53444F4D, // SDOM
+	kVersElement = 0x56455253, // VERS
+	kBordElement = 0x424F5244, // BORD
+	kSepoElement = 0x5345504F, // SEPO
+	kEcidElement = 0x45434944, // ECID
+	kUnknElement = 0x53414c54  // FIXME
 } img3_element_type;
 
 typedef struct {
-    unsigned int signature;
-    unsigned int full_size;
-    unsigned int data_size;
-    unsigned int shsh_offset;
-    unsigned int image_type;
+	unsigned int signature;
+	unsigned int full_size;
+	unsigned int data_size;
+	unsigned int shsh_offset;
+	unsigned int image_type;
 } img3_header;
 
 typedef struct {
-    unsigned int signature;
-    unsigned int full_size;
-    unsigned int data_size;
+	unsigned int signature;
+	unsigned int full_size;
+	unsigned int data_size;
 } img3_element_header;
 
 typedef struct {
@@ -77,13 +82,19 @@ typedef struct {
 	img3_element* ecid_element;
 	img3_element* shsh_element;
 	img3_element* cert_element;
+	img3_element* unkn_element;
 } img3_file;
 
-img3_file* img3_parse_file(char* data, int size);
-img3_element* img3_parse_element(char* data);
-int img3_replace_signature(img3_file* image, char* signature);
 void img3_free(img3_file* image);
-int img3_get_data(img3_file* image, char** pdata, int* psize);
+img3_element* img3_parse_element(char* data);
 void img3_free_element(img3_element* element);
+img3_file* img3_parse_file(char* data, int size);
+int img3_get_data(img3_file* image, char** pdata, int* psize);
+int img3_replace_signature(img3_file* image, char* signature);
+
+
+#ifdef __cplusplus
+}s
+#endif
 
 #endif
diff --git a/src/ipsw.c b/src/ipsw.c
index ca28596..f08e2fd 100644
--- a/src/ipsw.c
+++ b/src/ipsw.c
@@ -20,6 +20,7 @@
  */
 
 #include <zip.h>
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
@@ -94,8 +95,10 @@ int ipsw_extract_to_file(const char* ipsw, const char* infile, const char* outfi
 
 	int i = 0;
 	int size = 0;
+	int bytes = 0;
 	int count = 0;
-	for (i = zstat.size; i > 0; i -= count) {
+	double progress = 0;
+	for(i = zstat.size; i > 0; i -= count) {
 		if (i < BUFSIZE)
 			size = i;
 		else
@@ -108,9 +111,11 @@ int ipsw_extract_to_file(const char* ipsw, const char* infile, const char* outfi
 			return -1;
 		}
 		fwrite(buffer, 1, count, fd);
-		debug(".");
+
+		bytes += size;
+		progress = ((double) bytes/ (double) zstat.size) * 100.0;
+		print_progress_bar(progress);
 	}
-	debug("\n");
 
 	fclose(fd);
 	zip_fclose(zfile);
@@ -119,7 +124,7 @@ int ipsw_extract_to_file(const char* ipsw, const char* infile, const char* outfi
 	return 0;
 }
 
-int ipsw_extract_to_memory(const char* ipsw, const char* infile, char** pbuffer, int* psize) {
+int ipsw_extract_to_memory(const char* ipsw, const char* infile, char** pbuffer, uint32_t* psize) {
 	ipsw_archive* archive = ipsw_open(ipsw);
 	if (archive == NULL || archive->zip == NULL) {
 		error("ERROR: Invalid archive\n");
diff --git a/src/ipsw.h b/src/ipsw.h
index f764611..cd11406 100644
--- a/src/ipsw.h
+++ b/src/ipsw.h
@@ -19,10 +19,15 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef IPSW_H
-#define IPSW_H
+#ifndef IDEVICERESTORE_IPSW_H
+#define IDEVICERESTORE_IPSW_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 #include <zip.h>
+#include <stdint.h>
 
 typedef struct {
 	int index;
@@ -31,7 +36,11 @@ typedef struct {
 	unsigned char* data;
 } ipsw_file;
 
-int ipsw_extract_to_memory(const char* ipsw, const char* infile, char** pbuffer, int* psize);
+int ipsw_extract_to_memory(const char* ipsw, const char* infile, char** pbuffer, uint32_t* psize);
 void ipsw_free_file(ipsw_file* file);
 
+#ifdef __cplusplus
+}
+#endif
+
 #endif
diff --git a/src/normal.c b/src/normal.c
index c7baefd..c9a1b45 100644
--- a/src/normal.c
+++ b/src/normal.c
@@ -19,10 +19,318 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <libirecovery.h>
+#include <libimobiledevice/lockdown.h>
+#include <libimobiledevice/libimobiledevice.h>
 
+#include "common.h"
 #include "normal.h"
+#include "recovery.h"
 
-int normal_get_ecid(uint64_t* ecid) {
+static int normal_device_connected = 0;
+
+void normal_device_callback(const idevice_event_t* event, void* userdata) {
+	struct idevicerestore_client_t* client = (struct idevicerestore_client_t*) userdata;
+	if (event->event == IDEVICE_DEVICE_ADD) {
+		normal_device_connected = 1;
+
+	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
+		normal_device_connected = 0;
+		client->flags &= FLAG_QUIT;
+	}
+}
+
+int normal_client_new(struct idevicerestore_client_t* client) {
+	struct normal_client_t* normal = (struct normal_client_t*) malloc(sizeof(struct normal_client_t));
+	if (normal == NULL) {
+		error("ERROR: Out of memory\n");
+		return -1;
+	}
+
+	if (normal_open_with_timeout(client) < 0) {
+		normal_client_free(client);
+		return -1;
+	}
+
+	client->normal = normal;
+	return 0;
+}
+
+void normal_client_free(struct idevicerestore_client_t* client) {
+	struct normal_client_t* normal = NULL;
+	if (client) {
+		normal = client->normal;
+		if(normal) {
+			if(normal->client) {
+				lockdownd_client_free(normal->client);
+				normal->client = NULL;
+			}
+			if(normal->device) {
+				idevice_free(normal->device);
+				normal->device = NULL;
+			}
+		}
+		free(normal);
+		client->normal = NULL;
+	}
+}
+
+int normal_check_mode(const char* uuid) {
+	char* type = NULL;
+	idevice_t device = NULL;
+	lockdownd_client_t lockdown = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	lockdownd_error_t lockdown_error = IDEVICE_E_SUCCESS;
+
+	device_error = idevice_new(&device, uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	lockdown_error = lockdownd_client_new(device, &lockdown, "idevicerestore");
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		idevice_free(device);
+		return -1;
+	}
+
+	lockdown_error = lockdownd_query_type(lockdown, &type);
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		lockdownd_client_free(lockdown);
+		idevice_free(device);
+		return -1;
+	}
+
+	lockdownd_client_free(lockdown);
+	idevice_free(device);
+	lockdown = NULL;
+	device = NULL;
+	return 0;
+}
+
+int normal_open_with_timeout(struct idevicerestore_client_t* client) {
+	int i = 0;
+	int attempts = 10;
+	idevice_t device = NULL;
+	lockdownd_client_t lockdownd = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	lockdownd_error_t lockdownd_error = LOCKDOWN_E_SUCCESS;
+
+	// no context exists so bail
+	if(client == NULL) {
+		return -1;
+	}
+
+	// create our normal client if it doesn't yet exist
+	if(client->normal == NULL) {
+		client->normal = (struct normal_client_t*) malloc(sizeof(struct normal_client_t));
+		if(client->normal == NULL) {
+			error("ERROR: Out of memory\n");
+			return -1;
+		}
+	}
+
+	device_error = idevice_event_subscribe(&normal_device_callback, NULL);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		error("ERROR: Unable to subscribe to device events\n");
+		return -1;
+	}
+
+	for (i = 1; i <= attempts; i++) {
+		if (normal_device_connected == 1) {
+			break;
+		}
+
+		if (i == attempts) {
+			error("ERROR: Unable to connect to device in normal mode\n");
+			return -1;
+		}
+
+		sleep(2);
+	}
+
+	device_error = idevice_new(&device, client->uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	lockdownd_error = lockdownd_client_new(device, &lockdownd, "idevicerestore");
+	if (lockdownd_error != LOCKDOWN_E_SUCCESS) {
+		//idevice_event_unsubscribe();
+		idevice_free(device);
+		return -1;
+	}
+
+	char* type = NULL;
+	lockdownd_error = lockdownd_query_type(lockdownd, &type);
+	if (lockdownd_error != LOCKDOWN_E_SUCCESS) {
+		lockdownd_client_free(lockdownd);
+		//idevice_event_unsubscribe();
+		idevice_free(device);
+		return -1;
+	}
+
+	client->normal->device = device;
+	client->normal->client = lockdownd;
+	return 0;
+}
+
+int normal_check_device(const char* uuid) {
+	int i = 0;
+	idevice_t device = NULL;
+	char* product_type = NULL;
+	plist_t product_type_node = NULL;
+	lockdownd_client_t lockdown = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	lockdownd_error_t lockdown_error = IDEVICE_E_SUCCESS;
+
+	device_error = idevice_new(&device, uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	lockdown_error = lockdownd_client_new_with_handshake(device, &lockdown, "idevicerestore");
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		idevice_free(device);
+		return -1;
+	}
+
+	lockdown_error = lockdownd_get_value(lockdown, NULL, "ProductType", &product_type_node);
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		lockdownd_client_free(lockdown);
+		idevice_free(device);
+		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);
+		return -1;
+	}
+	plist_get_string_val(product_type_node, &product_type);
+	plist_free(product_type_node);
+
+	for (i = 0; idevicerestore_devices[i].product != NULL; i++) {
+		if (!strcmp(product_type, idevicerestore_devices[i].product)) {
+			break;
+		}
+	}
+
+	return idevicerestore_devices[i].index;
+}
+
+int normal_enter_recovery(struct idevicerestore_client_t* client) {
+	idevice_t device = NULL;
+	irecv_client_t recovery = NULL;
+	lockdownd_client_t lockdown = NULL;
+	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	lockdownd_error_t lockdown_error = LOCKDOWN_E_SUCCESS;
+
+	device_error = idevice_new(&device, client->uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		error("ERROR: Unable to find device\n");
+		return -1;
+	}
+
+	lockdown_error = lockdownd_client_new(device, &lockdown, "idevicerestore");
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		error("ERROR: Unable to connect to lockdownd service\n");
+		idevice_free(device);
+		return -1;
+	}
+
+	lockdown_error = lockdownd_enter_recovery(lockdown);
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		error("ERROR: Unable to place device in recovery mode\n");
+		lockdownd_client_free(lockdown);
+		idevice_free(device);
+		return -1;
+	}
+
+	lockdownd_client_free(lockdown);
+	idevice_free(device);
+	lockdown = NULL;
+	device = NULL;
+
+	if (recovery_open_with_timeout(client) < 0) {
+		error("ERROR: Unable to enter recovery mode\n");
+		return -1;
+	}
+
+	recovery_error = irecv_send_command(recovery, "setenv auto-boot true");
+	if (recovery_error != IRECV_E_SUCCESS) {
+		error("ERROR: Unable to reset auto-boot variable\n");
+		irecv_close(recovery);
+		return -1;
+	}
+
+	recovery_error = irecv_send_command(recovery, "saveenv");
+	if (recovery_error != IRECV_E_SUCCESS) {
+		error("ERROR: Unable to save auto-boot variable\n");
+		irecv_close(recovery);
+		return -1;
+	}
+
+	//client->mode = &idevicerestore_modes[MODE_RECOVERY];
+	irecv_close(recovery);
+	recovery = NULL;
+	return 0;
+}
+
+int normal_get_cpid(const char* uuid, uint32_t* cpid) {
+	return 0;
+}
+
+int normal_get_bdid(const char* uuid, uint32_t* bdid) {
+	return 0;
+}
+
+int normal_get_ecid(const char* uuid, uint64_t* ecid) {
+	idevice_t device = NULL;
+	plist_t unique_chip_node = NULL;
+	lockdownd_client_t lockdown = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	lockdownd_error_t lockdown_error = IDEVICE_E_SUCCESS;
+
+	device_error = idevice_new(&device, uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	lockdown_error = lockdownd_client_new_with_handshake(device, &lockdown, "idevicerestore");
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		error("ERROR: Unable to connect to lockdownd\n");
+		idevice_free(device);
+		return -1;
+	}
+
+	lockdown_error = lockdownd_get_value(lockdown, NULL, "UniqueChipID", &unique_chip_node);
+	if (lockdown_error != LOCKDOWN_E_SUCCESS) {
+		error("ERROR: Unable to get UniqueChipID from lockdownd\n");
+		lockdownd_client_free(lockdown);
+		idevice_free(device);
+		return -1;
+	}
+
+	if (!unique_chip_node || plist_get_node_type(unique_chip_node) != PLIST_UINT) {
+		error("ERROR: Unable to get ECID\n");
+		lockdownd_client_free(lockdown);
+		idevice_free(device);
+		return -1;
+	}
+	plist_get_uint_val(unique_chip_node, ecid);
+	plist_free(unique_chip_node);
+
+	lockdownd_client_free(lockdown);
+	idevice_free(device);
+	lockdown = NULL;
+	device = NULL;
 	return 0;
 }
diff --git a/src/normal.h b/src/normal.h
index 3e2868d..e86bf14 100644
--- a/src/normal.h
+++ b/src/normal.h
@@ -19,11 +19,37 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef NORMAL_H
-#define NORMAL_H
+#ifndef IDEVICERESTORE_NORMAL_H
+#define IDEVICERESTORE_NORMAL_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 #include <stdint.h>
+#include <libimobiledevice/lockdown.h>
+#include <libimobiledevice/libimobiledevice.h>
+
+struct normal_client_t {
+	idevice_t device;
+	lockdownd_client_t client;
+	const char* ipsw;
+	plist_t tss;
+};
 
-int normal_get_ecid(uint64_t* ecid);
+
+int normal_check_mode(const char* uuid);
+int normal_check_device(const char* uuid);
+int normal_client_new(struct idevicerestore_client_t* client);
+void normal_client_free(struct idevicerestore_client_t* client);
+int normal_open_with_timeout(struct idevicerestore_client_t* client);
+int normal_enter_recovery(struct idevicerestore_client_t* client);
+int normal_get_cpid(const char* uuid, uint32_t* cpid);
+int normal_get_bdid(const char* uuid, uint32_t* cpid);
+int normal_get_ecid(const char* uuid, uint64_t* ecid);
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif
diff --git a/src/recovery.c b/src/recovery.c
index 4e2e7ad..6a38343 100644
--- a/src/recovery.c
+++ b/src/recovery.c
@@ -21,249 +21,389 @@
 
 #include <stdio.h>
 #include <stdlib.h>
-#include <stdint.h>
 #include <libirecovery.h>
+#include <libimobiledevice/restore.h>
+#include <libimobiledevice/libimobiledevice.h>
 
 #include "tss.h"
 #include "img3.h"
+#include "common.h"
+#include "restore.h"
 #include "recovery.h"
 #include "idevicerestore.h"
 
-int recovery_send_signed_component(irecv_client_t client, char* ipsw, plist_t tss, char* component) {
-	int size = 0;
-	char* data = NULL;
-	char* path = NULL;
-	char* blob = NULL;
-	img3_file* img3 = NULL;
-	irecv_error_t error = 0;
-
-	if (get_signed_component_by_name(ipsw, tss, component, &data, &size) < 0) {
-		error("ERROR: Unable to get signed component: %s\n", component);
-		return -1;
+int recovery_progress_callback(irecv_client_t client, const irecv_event_t* event) {
+	if (event->type == IRECV_PROGRESS) {
+		print_progress_bar(event->progress);
 	}
+	return 0;
+}
 
-	info("Sending %s...\n", component);
-	error = irecv_send_buffer(client, data, size);
-	if (error != IRECV_E_SUCCESS) {
-		error("ERROR: Unable to send IMG3: %s\n", path);
-		img3_free(img3);
-		free(data);
-		free(path);
+int recovery_client_new(struct idevicerestore_client_t* client) {
+	struct recovery_client_t* recovery = (struct recovery_client_t*) malloc(sizeof(struct recovery_client_t));
+	if (recovery == NULL) {
+		error("ERROR: Out of memory\n");
 		return -1;
 	}
 
-	if (data) {
-		free(data);
-		data = NULL;
+	client->recovery = recovery;
+
+	if (recovery_open_with_timeout(client) < 0) {
+		recovery_client_free(client);
+		return -1;
 	}
 
+	client->recovery = recovery;
 	return 0;
 }
 
-irecv_error_t recovery_open_with_timeout(irecv_client_t* client) {
+void recovery_client_free(struct idevicerestore_client_t* client) {
+	if(client) {
+		if (client->recovery) {
+			if(client->recovery->client) {
+				irecv_close(client->recovery->client);
+				client->recovery->client = NULL;
+			}
+			free(client->recovery);
+			client->recovery = NULL;
+		}
+	}
+}
+
+int recovery_open_with_timeout(struct idevicerestore_client_t* client) {
 	int i = 0;
-	irecv_error_t error = 0;
-	for (i = 10; i > 0; i--) {
-		error = irecv_open(client);
-		if (error == IRECV_E_SUCCESS) {
-			return error;
+	int attempts = 10;
+	irecv_client_t recovery = NULL;
+	irecv_error_t recovery_error = IRECV_E_UNKNOWN_ERROR;
+
+	if(client->recovery == NULL) {
+		if(recovery_client_new(client) < 0) {
+			error("ERROR: Unable to open device in recovery mode\n");
+			return -1;
+		}
+		return 0;
+	}
+
+	for (i = 1; i <= attempts; i++) {
+		recovery_error = irecv_open(&recovery);
+		if (recovery_error == IRECV_E_SUCCESS) {
+			break;
+		}
+
+		if (i >= attempts) {
+			error("ERROR: Unable to connect to device in recovery mode\n");
+			return -1;
 		}
 
 		sleep(2);
-		info("Retrying connection...\n");
+		debug("Retrying connection...\n");
 	}
 
-	error("ERROR: Unable to connect to recovery device.\n");
-	return error;
+	irecv_event_subscribe(recovery, IRECV_PROGRESS, &recovery_progress_callback, NULL);
+	client->recovery->client = recovery;
+	return 0;
 }
 
-int recovery_send_ibec(char* ipsw, plist_t tss) {
+int recovery_check_mode() {
+	irecv_client_t recovery = NULL;
+	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+
+	recovery_error = irecv_open(&recovery);
+	if (recovery_error != IRECV_E_SUCCESS) {
+		return -1;
+	}
+
+	if (recovery->mode == kDfuMode) {
+		irecv_close(recovery);
+		return -1;
+	}
+
+	irecv_close(recovery);
+	recovery = NULL;
+	return 0;
+}
+
+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, build_identity) < 0) {
+		error("ERROR: Unable to send iBEC\n");
+		return -1;
+	}
+	sleep(2);
+
+	if (recovery_send_applelogo(client, build_identity) < 0) {
+		error("ERROR: Unable to send AppleLogo\n");
+		return -1;
+	}
+
+	if (recovery_send_devicetree(client, build_identity) < 0) {
+		error("ERROR: Unable to send DeviceTree\n");
+		return -1;
+	}
+
+	if (recovery_send_ramdisk(client, build_identity) < 0) {
+		error("ERROR: Unable to send Ramdisk\n");
+		return -1;
+	}
+
+	// for some reason iboot requires a hard reset after ramdisk
+	//  or things start getting wacky
+	printf("Please unplug your device, then plug it back in\n");
+	printf("Hit any key to continue...");
+	getchar();
+
+	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;
+	}
+
+	info("Waiting for device to enter restore mode\n");
+	if (restore_open_with_timeout(client) < 0) {
+		error("ERROR: Unable to connect to device in restore mode\n");
+		return -1;
+	}
+
+	restore_client_free(client);
+	client->mode = &idevicerestore_modes[MODE_RESTORE];
+	return 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;
-	irecv_client_t client = NULL;
-	char* component = "iBEC";
 
-	error = recovery_open_with_timeout(&client);
-	if (error != IRECV_E_SUCCESS) {
+	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;
+		}
+	}
+
+	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;
 	}
 
-	error = irecv_send_command(client, "setenv auto-boot true");
+	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 %s component: %s\n", component, irecv_strerror(error));
+		free(data);
+		return -1;
+	}
+
+	free(data);
+	return 0;
+}
+
+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) {
 		error("ERROR: Unable to set auto-boot environmental variable\n");
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "saveenv");
-	if (error != IRECV_E_SUCCESS) {
+	recovery_error = irecv_send_command(client->recovery->client, "saveenv");
+	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to save environmental variable\n");
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, ipsw, tss, component) < 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);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "go");
-	if (error != IRECV_E_SUCCESS) {
+	recovery_error = irecv_send_command(client->recovery->client, "go");
+	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to execute %s\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	if (client) {
-		irecv_close(client);
-		client = NULL;
-	}
 	return 0;
 }
 
-int recovery_send_applelogo(char* ipsw, plist_t tss) {
-	irecv_error_t error = 0;
-	irecv_client_t client = NULL;
-	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);
-
-	error = recovery_open_with_timeout(&client);
-	if (error != IRECV_E_SUCCESS) {
+	if (recovery_open_with_timeout(client) < 0) {
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, ipsw, tss, component) < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "setpicture 1");
-	if (error != IRECV_E_SUCCESS) {
+	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(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "bgcolor 0 0 0");
-	if (error != IRECV_E_SUCCESS) {
+	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(client);
-		client = NULL;
 		return -1;
 	}
 
-	if (client) {
-		irecv_close(client);
-		client = NULL;
-	}
 	return 0;
 }
 
-int recovery_send_devicetree(char* ipsw, plist_t tss) {
-	irecv_error_t error = 0;
-	irecv_client_t client = NULL;
-	char *component = "RestoreDeviceTree";
+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;
 
-	error = recovery_open_with_timeout(&client);
-	if (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, ipsw, tss, component) < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "devicetree");
-	if (error != IRECV_E_SUCCESS) {
+	recovery_error = irecv_send_command(client->recovery->client, "devicetree");
+	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to execute %s\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	if (client) {
-		irecv_close(client);
-		client = NULL;
-	}
 	return 0;
 }
 
-int recovery_send_ramdisk(char* ipsw, plist_t tss) {
-	irecv_error_t error = 0;
-	irecv_client_t client = NULL;
-	char *component = "RestoreRamDisk";
+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;
 
-	error = recovery_open_with_timeout(&client);
-	if (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, ipsw, tss, component) < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "ramdisk");
-	if (error != IRECV_E_SUCCESS) {
+	recovery_error = irecv_send_command(client->recovery->client, "ramdisk");
+	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to execute %s\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	if (client) {
-		irecv_close(client);
-		client = NULL;
-	}
 	return 0;
 }
 
-int recovery_send_kernelcache(char* ipsw, plist_t tss) {
-	irecv_error_t error = 0;
-	irecv_client_t client = NULL;
-	char *component = "RestoreKernelCache";
+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;
 
-	error = recovery_open_with_timeout(&client);
-	if (error != IRECV_E_SUCCESS) {
+	if (recovery_open_with_timeout(client) < 0) {
 		return -1;
 	}
 
-	if (recovery_send_signed_component(client, ipsw, tss, component) < 0) {
+	if (recovery_send_component(client, build_identity, component) < 0) {
 		error("ERROR: Unable to send %s to device.\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	error = irecv_send_command(client, "bootx");
-	if (error != IRECV_E_SUCCESS) {
+	recovery_error = irecv_send_command(client->recovery->client, "bootx");
+	if (recovery_error != IRECV_E_SUCCESS) {
 		error("ERROR: Unable to execute %s\n", component);
-		irecv_close(client);
-		client = NULL;
 		return -1;
 	}
 
-	if (client) {
-		irecv_close(client);
-		client = NULL;
+	return 0;
+}
+
+int recovery_get_ecid(struct idevicerestore_client_t* client, uint64_t* ecid) {
+	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+
+	if(client->recovery == NULL) {
+		if (recovery_open_with_timeout(client) < 0) {
+			return -1;
+		}
 	}
+
+	recovery_error = irecv_get_ecid(client->recovery->client, ecid);
+	if (recovery_error != IRECV_E_SUCCESS) {
+		return -1;
+	}
+
 	return 0;
 }
 
+int recovery_get_cpid(struct idevicerestore_client_t* client, uint32_t* cpid) {
+	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+
+	if(client->recovery == NULL) {
+		if (recovery_open_with_timeout(client) < 0) {
+			return -1;
+		}
+	}
+
+	recovery_error = irecv_get_cpid(client->recovery->client, cpid);
+	if (recovery_error != IRECV_E_SUCCESS) {
+		return -1;
+	}
+
+	return 0;
+}
+
+int recovery_get_bdid(struct idevicerestore_client_t* client, uint32_t* bdid) {
+	irecv_error_t recovery_error = IRECV_E_SUCCESS;
+
+	if(client->recovery == NULL) {
+		if (recovery_open_with_timeout(client) < 0) {
+			return -1;
+		}
+	}
+
+	recovery_error = irecv_get_bdid(client->recovery->client, bdid);
+	if (recovery_error != IRECV_E_SUCCESS) {
+		return -1;
+	}
 
-int recovery_get_ecid(uint64_t* ecid) {
 	return 0;
 }
diff --git a/src/recovery.h b/src/recovery.h
index 5495638..b7cc0e4 100644
--- a/src/recovery.h
+++ b/src/recovery.h
@@ -19,19 +19,43 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef RECOVERY_H
-#define RECOVERY_H
+#ifndef IDEVICERESTORE_RECOVERY_H
+#define IDEVICERESTORE_RECOVERY_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 #include <stdint.h>
 #include <plist/plist.h>
 
-int recovery_send_signed_component(irecv_client_t client, char* ipsw, plist_t tss, char* component);
-irecv_error_t recovery_open_with_timeout(irecv_client_t* client);
-int recovery_send_ibec(char* ipsw, plist_t tss);
-int recovery_send_applelogo(char* ipsw, plist_t tss);
-int recovery_send_devicetree(char* ipsw, plist_t tss);
-int recovery_send_ramdisk(char* ipsw, plist_t tss);
-int recovery_send_kernelcache(char* ipsw, plist_t tss);
-int recovery_get_ecid(uint64_t* ecid);
+#include "common.h"
+
+struct irecv_client;
+typedef struct irecv_client* irecv_client_t;
+struct recovery_client_t {
+	irecv_client_t client;
+	const char* ipsw;
+	plist_t tss;
+};
+
+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_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);
+
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif
diff --git a/src/restore.c b/src/restore.c
index 485df9b..5d45296 100644
--- a/src/restore.c
+++ b/src/restore.c
@@ -24,10 +24,10 @@
 #include <string.h>
 #include <libimobiledevice/restore.h>
 
+#include "asr.h"
+#include "tss.h"
+#include "common.h"
 #include "restore.h"
-#include "idevicerestore.h"
-
-#define ASR_PORT 12345
 
 #define CREATE_PARTITION_MAP   12
 #define CREATE_FILESYSTEM      13
@@ -44,8 +44,233 @@
 #define WAIT_FOR_DEVICE        33
 #define LOAD_NOR               36
 
+static int restore_device_connected = 0;
+
+int restore_client_new(struct idevicerestore_client_t* client) {
+	struct restore_client_t* restore = (struct restore_client_t*) malloc(sizeof(struct restore_client_t));
+	if (restore == NULL) {
+		error("ERROR: Out of memory\n");
+		return -1;
+	}
+
+	if (restore_open_with_timeout(client) < 0) {
+		restore_client_free(client);
+		return -1;
+	}
+
+	client->restore = restore;
+	return 0;
+}
+
+void restore_client_free(struct idevicerestore_client_t* client) {
+	if (client) {
+		if(client->restore) {
+			if(client->restore->client) {
+				restored_client_free(client->restore->client);
+				client->restore->client = NULL;
+			}
+			if(client->restore->device) {
+				idevice_free(client->restore->device);
+				client->restore->device = NULL;
+			}
+			free(client->restore);
+			client->restore = NULL;
+		}
+	}
+}
+
+int restore_check_mode(const char* uuid) {
+	char* type = 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;
+	}
+
+	restored_client_free(restore);
+	idevice_free(device);
+	restore = NULL;
+	device = NULL;
+	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].index;
+}
+
+void restore_device_callback(const idevice_event_t* event, void* userdata) {
+	struct idevicerestore_client_t* client = (struct idevicerestore_client_t*) userdata;
+	if (event->event == IDEVICE_DEVICE_ADD) {
+		restore_device_connected = 1;
+
+	} else if (event->event == IDEVICE_DEVICE_REMOVE) {
+		restore_device_connected = 0;
+		client->flags &= FLAG_QUIT;
+	}
+}
+
+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 == NULL) {
+		if (restore_open_with_timeout(client) < 0) {
+			error("ERROR: Unable to open device in restore mode\n");
+			return -1;
+		}
+	}
+
+	restore_error = restored_reboot(client->restore->client);
+	if (restore_error != RESTORE_E_SUCCESS) {
+		error("ERROR: Unable to reboot the device from restore mode\n");
+		return -1;
+	}
+
+	return 0;
+}
+
+int restore_open_with_timeout(struct idevicerestore_client_t* client) {
+	int i = 0;
+	int attempts = 10;
+	idevice_t device = NULL;
+	restored_client_t restored = NULL;
+	idevice_error_t device_error = IDEVICE_E_SUCCESS;
+	restored_error_t restore_error = RESTORE_E_SUCCESS;
+
+	// no context exists so bail
+	if(client == NULL) {
+		return -1;
+	}
+
+	// create our restore client if it doesn't yet exist
+	if(client->restore == NULL) {
+		client->restore = (struct restore_client_t*) malloc(sizeof(struct restore_client_t));
+		if(client->restore == NULL) {
+			error("ERROR: Out of memory\n");
+			return -1;
+		}
+	}
+
+	device_error = idevice_event_subscribe(&restore_device_callback, NULL);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		error("ERROR: Unable to subscribe to device events\n");
+		return -1;
+	}
+
+	for (i = 1; i <= attempts; i++) {
+		if (restore_device_connected == 1) {
+			break;
+		}
+
+		if (i == attempts) {
+			error("ERROR: Unable to connect to device in restore mode\n");
+			return -1;
+		}
+
+		sleep(2);
+	}
+
+	device_error = idevice_new(&device, client->uuid);
+	if (device_error != IDEVICE_E_SUCCESS) {
+		return -1;
+	}
+
+	restore_error = restored_client_new(device, &restored, "idevicerestore");
+	if (restore_error != RESTORE_E_SUCCESS) {
+		//idevice_event_unsubscribe();
+		idevice_free(device);
+		return -1;
+	}
+
+	restore_error = restored_query_type(restored, NULL, NULL);
+	if (restore_error != RESTORE_E_SUCCESS) {
+		restored_client_free(restored);
+		//idevice_event_unsubscribe();
+		idevice_free(device);
+		return -1;
+	}
+
+	client->restore->device = device;
+	client->restore->client = restored;
+	return 0;
+}
+
 const char* restore_progress_string(unsigned int operation) {
-	switch(operation) {
+	switch (operation) {
 	case CREATE_PARTITION_MAP:
 		return "Creating partition map";
 
@@ -93,217 +318,122 @@ const char* restore_progress_string(unsigned int operation) {
 	}
 }
 
-
 int restore_handle_progress_msg(restored_client_t client, plist_t msg) {
 	plist_t node = NULL;
+	uint64_t progress = 0;
 	uint64_t operation = 0;
-	uint64_t uprogress = 0;
-	uint32_t progress = 0;
 
 	node = plist_dict_get_item(msg, "Operation");
-	if (node && PLIST_UINT == plist_get_node_type(node)) {
-		plist_get_uint_val(node, &operation);
-	} else {
+	if (!node || plist_get_node_type(node) != PLIST_UINT) {
 		debug("Failed to parse operation from ProgressMsg plist\n");
-		return 0;
+		return -1;
 	}
+	plist_get_uint_val(node, &operation);
 
 	node = plist_dict_get_item(msg, "Progress");
-	if (node && PLIST_UINT == plist_get_node_type(node)) {
-		plist_get_uint_val(node, &uprogress);
-		progress = (uint32_t) uprogress;
-	} else {
+	if (!node || plist_get_node_type(node) != PLIST_UINT) {
 		debug("Failed to parse progress from ProgressMsg plist \n");
-		return 0;
+		return -1;
 	}
+	plist_get_uint_val(node, &progress);
+
+	if ((progress > 0) && (progress < 100)) {
+		print_progress_bar((double) progress);
 
-	if ((progress > 0) && (progress < 100))
-		info("%s - Progress: %ul%%\n", restore_progress_string(operation), progress);
-	else
+	} else {
 		info("%s\n", restore_progress_string(operation));
+	}
 
 	return 0;
 }
 
 int restore_handle_status_msg(restored_client_t client, plist_t msg) {
 	info("Got status message\n");
+	debug_plist(msg);
 	return 0;
 }
 
-int asr_send_system_image_data_from_file(idevice_t device, restored_client_t client, const char *filesystem) {
+int restore_send_filesystem(idevice_t device, const char* filesystem) {
 	int i = 0;
-	char buffer[0x1000];
-	uint32_t recv_bytes = 0;
-	memset(buffer, '\0', 0x1000);
-	idevice_connection_t connection = NULL;
-	idevice_error_t ret = IDEVICE_E_UNKNOWN_ERROR;
-
-	for (i = 0; i < 5; i++) {
-		ret = idevice_connect(device, ASR_PORT, &connection);
-		if (ret == IDEVICE_E_SUCCESS)
-			break;
+	FILE* file = NULL;
+	plist_t data = NULL;
+	idevice_connection_t asr = NULL;
+	idevice_error_t device_error = IDEVICE_E_UNKNOWN_ERROR;
 
-		else
-			sleep(1);
+	if (asr_open_with_timeout(device, &asr) < 0) {
+		error("ERROR: Unable to connect to ASR\n");
+		return -1;
 	}
+	info("Connected to ASR\n");
 
-	if (ret != IDEVICE_E_SUCCESS)
-		return ret;
-
-	memset(buffer, '\0', 0x1000);
-	ret = idevice_connection_receive(connection, buffer, 0x1000, &recv_bytes);
-	if (ret != IDEVICE_E_SUCCESS) {
-		idevice_disconnect(connection);
-		return ret;
+	// we don't really need to do anything with this,
+	// we're just clearing the output buffer
+	if (asr_receive(asr, &data) < 0) {
+		error("ERROR: Unable to receive data from ASR\n");
+		asr_close(asr);
+		return -1;
 	}
-	info("Received %d bytes\n", recv_bytes);
-	info("%s", buffer);
-
-	FILE* fd = fopen(filesystem, "rb");
-	if (fd == NULL) {
-		idevice_disconnect(connection);
-		return ret;
+	plist_free(data);
+
+	// this step sends requested chunks of data from various offsets to asr so
+	// it can validate the filesystem before installing it
+	debug("Preparing to validate the filesystem\n");
+	if (asr_perform_validation(asr, filesystem) < 0) {
+		error("ERROR: ASR was unable to validate the filesystem\n");
+		asr_close(asr);
+		return -1;
 	}
-
-	fseek(fd, 0, SEEK_END);
-	uint64_t len = ftell(fd);
-	fseek(fd, 0, SEEK_SET);
-
-	info("Connected to ASR\n");
-	plist_t dict = plist_new_dict();
-	plist_dict_insert_item(dict, "FEC Slice Stride", plist_new_uint(40));
-	plist_dict_insert_item(dict, "Packet Payload Size", plist_new_uint(1450));
-	plist_dict_insert_item(dict, "Packets Per FEC", plist_new_uint(25));
-
-	plist_t payload = plist_new_dict();
-	plist_dict_insert_item(payload, "Port", plist_new_uint(1));
-	plist_dict_insert_item(payload, "Size", plist_new_uint(len));
-	plist_dict_insert_item(dict, "Payload", payload);
-
-	plist_dict_insert_item(dict, "Stream ID", plist_new_uint(1));
-	plist_dict_insert_item(dict, "Version", plist_new_uint(1));
-
-	char* xml = NULL;
-	unsigned int dict_size = 0;
-	unsigned int sent_bytes = 0;
-	plist_to_xml(dict, &xml, &dict_size);
-
-	ret = idevice_connection_send(connection, xml, dict_size, &sent_bytes);
-	if (ret != IDEVICE_E_SUCCESS) {
-		idevice_disconnect(connection);
-		return ret;
+	info("Filesystem validated\n");
+
+	// once the target filesystem has been validated, ASR then requests the
+	// entire filesystem to be sent.
+	debug("Preparing to send filesystem\n");
+	if (asr_send_payload(asr, filesystem) < 0) {
+		error("ERROR: Unable to send payload to ASR\n");
+		asr_close(asr);
+		return -1;
 	}
+	info("Filesystem finished\n");
 
-	info("Sent %d bytes\n", sent_bytes);
-	info("%s", xml);
-	plist_free(dict);
-	free(xml);
-
-	char* command = NULL;
-	do {
-		memset(buffer, '\0', 0x1000);
-		ret = idevice_connection_receive(connection, buffer, 0x1000, &recv_bytes);
-		if (ret != IDEVICE_E_SUCCESS) {
-			idevice_disconnect(connection);
-			return ret;
-		}
-		info("Received %d bytes\n", recv_bytes);
-		info("%s", buffer);
-
-		plist_t request = NULL;
-		plist_from_xml(buffer, recv_bytes, &request);
-		plist_t command_node = plist_dict_get_item(request, "Command");
-		if (command_node && PLIST_STRING == plist_get_node_type(command_node)) {
-			plist_get_string_val(command_node, &command);
-			if (!strcmp(command, "OOBData")) {
-				plist_t oob_length_node = plist_dict_get_item(request, "OOB Length");
-				if (!oob_length_node || PLIST_UINT != plist_get_node_type(oob_length_node)) {
-					error("Error fetching OOB Length\n");
-					idevice_disconnect(connection);
-					return IDEVICE_E_UNKNOWN_ERROR;
-				}
-				uint64_t oob_length = 0;
-				plist_get_uint_val(oob_length_node, &oob_length);
-
-				plist_t oob_offset_node = plist_dict_get_item(request, "OOB Offset");
-				if (!oob_offset_node || PLIST_UINT != plist_get_node_type(oob_offset_node)) {
-					error("Error fetching OOB Offset\n");
-					idevice_disconnect(connection);
-					return IDEVICE_E_UNKNOWN_ERROR;
-				}
-				uint64_t oob_offset = 0;
-				plist_get_uint_val(oob_offset_node, &oob_offset);
-
-				char* oob_data = (char*) malloc(oob_length);
-				if (oob_data == NULL) {
-					error("Out of memory\n");
-					idevice_disconnect(connection);
-					return IDEVICE_E_UNKNOWN_ERROR;
-				}
-
-				fseek(fd, oob_offset, SEEK_SET);
-				if (fread(oob_data, 1, oob_length, fd) != oob_length) {
-					error("Unable to read filesystem offset\n");
-					idevice_disconnect(connection);
-					free(oob_data);
-					return ret;
-				}
-
-				ret = idevice_connection_send(connection, oob_data, oob_length, &sent_bytes);
-				if (sent_bytes != oob_length || ret != IDEVICE_E_SUCCESS) {
-					error("Unable to send %d bytes to asr\n", sent_bytes);
-					idevice_disconnect(connection);
-					free(oob_data);
-					return ret;
-				}
-				plist_free(request);
-				free(oob_data);
-			}
-		}
-
-	} while (strcmp(command, "Payload"));
+	asr_close(asr);
+	return 0;
+}
 
-	fseek(fd, 0, SEEK_SET);
-	char data[1450];
-	for (i = len; i > 0; i -= 1450) {
-		int size = 1450;
-		if (i < 1450) {
-			size = i;
-		}
+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;
+	plist_t blob = NULL;
+	plist_t dict = NULL;
+	restored_error_t restore_error = RESTORE_E_SUCCESS;
 
-		if (fread(data, 1, size, fd) != (unsigned int) size) {
-			fclose(fd);
-			idevice_disconnect(connection);
-			error("Error reading filesystem\n");
-			return IDEVICE_E_UNKNOWN_ERROR;
-		}
+	info("Sending kernelcache\n");
 
-		ret = idevice_connection_send(connection, data, size, &sent_bytes);
-		if (ret != IDEVICE_E_SUCCESS) {
-			fclose(fd);
+	if (client->tss) {
+		if (tss_get_entry_path(client->tss, "KernelCache", &path) < 0) {
+			error("ERROR: Unable to get KernelCache path\n");
+			return -1;
 		}
-
-		if (i % (1450 * 1000) == 0) {
-			info(".");
+	} 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;
 		}
 	}
 
-	info("Done sending filesystem\n");
-	fclose(fd);
-	ret = idevice_disconnect(connection);
-	return ret;
-}
-
-int restore_send_kernelcache(restored_client_t client, char *kernel_data, int len) {
-	info("Sending kernelcache\n");
-
-	plist_t kernelcache_node = plist_new_data(kernel_data, len);
+	if (ipsw_get_component_by_path(client->ipsw, client->tss, path, &data, &size) < 0) {
+		error("ERROR: Unable to get kernelcache file\n");
+		return -1;
+	}
 
-	plist_t dict = plist_new_dict();
-	plist_dict_insert_item(dict, "KernelCacheFile", kernelcache_node);
+	dict = plist_new_dict();
+	blob = plist_new_data(data, size);
+	plist_dict_insert_item(dict, "KernelCacheFile", blob);
 
-	restored_error_t ret = restored_send(client, dict);
-	if (ret != RESTORE_E_SUCCESS) {
+	restore_error = restored_send(restore, dict);
+	if (restore_error != RESTORE_E_SUCCESS) {
 		error("ERROR: Unable to send kernelcache data\n");
 		plist_free(dict);
 		return -1;
@@ -311,55 +441,70 @@ int restore_send_kernelcache(restored_client_t client, char *kernel_data, int le
 
 	info("Done sending kernelcache\n");
 	plist_free(dict);
+	free(data);
 	return 0;
 }
 
-int restore_send_nor_data(restored_client_t client, 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_blob = NULL;
-	if (get_tss_data_by_name(tss, "LLB", &llb_path, &llb_blob) < 0) {
-		error("ERROR: Unable to get LLB info from TSS response\n");
-		return -1;
+	char* llb_filename = NULL;
+	char firmware_path[256];
+	char manifest_file[256];
+	int manifest_size = 0;
+	char* manifest_data = NULL;
+	char firmware_filename[256];
+	int llb_size = 0;
+	char* llb_data = NULL;
+	plist_t dict = NULL;
+	char* filename = NULL;
+	int nor_size = 0;
+	char* nor_data = NULL;
+	plist_t norimage_array = NULL;
+	restored_error_t ret = RESTORE_E_SUCCESS;
+
+	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;
+		}
 	}
 
-	char* llb_filename = strstr(llb_path, "LLB");
+	llb_filename = strstr(llb_path, "LLB");
 	if (llb_filename == NULL) {
-		error("ERROR: Unable to extrac firmware path from LLB filename\n");
+		error("ERROR: Unable to extract firmware path from LLB filename\n");
 		free(llb_path);
-		free(llb_blob);
 		return -1;
 	}
 
-	char firmware_path[256];
 	memset(firmware_path, '\0', sizeof(firmware_path));
 	memcpy(firmware_path, llb_path, (llb_filename - 1) - llb_path);
 	info("Found firmware path %s\n", firmware_path);
 
-	char manifest_file[256];
 	memset(manifest_file, '\0', sizeof(manifest_file));
 	snprintf(manifest_file, sizeof(manifest_file), "%s/manifest", firmware_path);
 	info("Getting firmware manifest %s\n", manifest_file);
 
-	int manifest_size = 0;
-	char* manifest_data = NULL;
-	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);
-		free(llb_blob);
 		return -1;
 	}
 
-	char firmware_filename[256];
 	memset(firmware_filename, '\0', sizeof(firmware_filename));
 
-	int llb_size = 0;
-	char* llb_data = NULL;
-	plist_t dict = plist_new_dict();
-	char* filename = strtok(manifest_data, "\n");
+	dict = plist_new_dict();
+	filename = strtok(manifest_data, "\n");
 	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_by_path(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;
 		}
@@ -367,14 +512,12 @@ int restore_send_nor_data(restored_client_t client, char* ipsw, plist_t tss) {
 		plist_dict_insert_item(dict, "LlbImageData", plist_new_data(llb_data, (uint64_t) llb_size));
 	}
 
-	int nor_size = 0;
-	char* nor_data = NULL;
 	filename = strtok(NULL, "\n");
-	plist_t norimage_array = plist_new_array();
+	norimage_array = plist_new_array();
 	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_by_path(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;
 		}
@@ -387,13 +530,9 @@ int restore_send_nor_data(restored_client_t client, char* ipsw, plist_t tss) {
 	}
 	plist_dict_insert_item(dict, "NorImageData", norimage_array);
 
-	int sz = 0;
-	char* xml = NULL;
-	plist_to_xml(dict, &xml, &sz);
-	debug("%s", xml);
-	free(xml);
+	debug_plist(dict);
 
-	restored_error_t 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);
@@ -403,3 +542,137 @@ int restore_send_nor_data(restored_client_t client, char* ipsw, plist_t tss) {
 	plist_free(dict);
 	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 build_identity, const char* filesystem) {
+	char* type = NULL;
+	plist_t node = NULL;
+
+	// checks and see what kind of data restored is requests and pass
+	// the request to its own handler
+	node = plist_dict_get_item(message, "DataType");
+	if (node && PLIST_STRING == plist_get_node_type(node)) {
+		plist_get_string_val(node, &type);
+
+		// this request is sent when restored is ready to receive the filesystem
+		if (!strcmp(type, "SystemImageData")) {
+			if(restore_send_filesystem(device, filesystem) < 0) {
+				error("ERROR: Unable to send filesystem\n");
+				return -1;
+			}
+		}
+
+		else if (!strcmp(type, "KernelCache")) {
+			if(restore_send_kernelcache(restore, client, build_identity) < 0) {
+				error("ERROR: Unable to send kernelcache\n");
+				return -1;
+			}
+		}
+
+		else if (!strcmp(type, "NORData")) {
+			if(client->flags & FLAG_EXCLUDE > 0) {
+				if(restore_send_nor(restore, client, build_identity) < 0) {
+					error("ERROR: Unable to send NOR data\n");
+					return -1;
+				}
+			} else {
+				client->flags &= 1;
+			}
+
+		} else {
+			// Unknown DataType!!
+			debug("Unknown data request received\n");
+			debug_plist(message);
+		}
+	}
+	return 0;
+}
+
+int restore_device(struct idevicerestore_client_t* client, plist_t build_identity, const char* filesystem) {
+	int error = 0;
+	char* type = NULL;
+	char* kernel = NULL;
+	plist_t node = NULL;
+	plist_t message = NULL;
+	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;
+
+	// open our connection to the device and verify we're in restore mode
+	if (restore_open_with_timeout(client) < 0) {
+		error("ERROR: Unable to open device in restore mode\n");
+		return -1;
+	}
+	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) {
+		error("ERROR: Unable to start the restore process\n");
+		restore_client_free(client);
+		return -1;
+	}
+
+	// this is the restore process loop, it reads each message in from
+	// restored and passes that data on to it's specific handler
+	while ((client->flags & FLAG_QUIT) == 0) {
+		restore_error = restored_receive(restore, &message);
+		if (restore_error != RESTORE_E_SUCCESS) {
+			debug("No data to read\n");
+			message = NULL;
+			continue;
+		}
+
+		// discover what kind of message has been received
+		node = plist_dict_get_item(message, "MsgType");
+		if (!node || plist_get_node_type(node) != PLIST_STRING) {
+			debug("Unknown message received\n");
+			debug_plist(message);
+			plist_free(message);
+			message = NULL;
+			continue;
+		}
+		plist_get_string_val(node, &type);
+
+		// data request messages are sent by restored whenever it requires
+		// 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, build_identity, filesystem);
+		}
+
+		// progress notification messages sent by the restored inform the client
+		// of it's current operation and sometimes percent of progress is complete
+		else if (!strcmp(type, "ProgressMsg")) {
+			error = restore_handle_progress_msg(restore, message);
+		}
+
+		// status messages usually indicate the current state of the restored
+		// process or often to signal an error has been encountered
+		else if (!strcmp(type, "StatusMsg")) {
+			error = restore_handle_status_msg(restore, message);
+		}
+
+		// there might be some other message types i'm not aware of, but I think
+		// at least the "previous error logs" messages usually end up here
+		else {
+			debug("Unknown message type received\n");
+			debug_plist(message);
+		}
+
+		// finally, if any of these message handlers returned -1 then we encountered
+		// an unrecoverable error, so we need to bail.
+		if (error < 0) {
+			error("ERROR: Unable to successfully restore device\n");
+			client->flags &= FLAG_QUIT;
+		}
+
+		plist_free(message);
+		message = NULL;
+	}
+
+	restore_client_free(client);
+	return 0;
+}
diff --git a/src/restore.h b/src/restore.h
index 644658a..9c11c34 100644
--- a/src/restore.h
+++ b/src/restore.h
@@ -1,4 +1,4 @@
-/*
+			/*
  * restore.h
  * Functions for handling idevices in restore mode
  *
@@ -19,18 +19,44 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef RESTORED_H
-#define RESTORED_H
+#ifndef IDEVICERESTORE_RESTORE_H
+#define IDEVICERESTORE_RESTORE_H
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <plist/plist.h>
 #include <libimobiledevice/restore.h>
+#include <libimobiledevice/libimobiledevice.h>
 
-#include "restore.h"
+struct restore_client_t {
+	plist_t tss;
+	idevice_t device;
+	const char* uuid;
+	unsigned int operation;
+	const char* filesystem;
+	restored_client_t client;
+};
 
-int restore_handle_progress_msg(restored_client_t client, plist_t msg);
-int restore_handle_status_msg(restored_client_t client, plist_t msg);
-int asr_send_system_image_data_from_file(idevice_t device, restored_client_t client, const char *filesystem);
-int restore_send_kernelcache(restored_client_t client, char *kernel_data, int len);
-int restore_send_nor_data(restored_client_t client, char* ipsw, plist_t tss);
+int restore_check_mode(const char* uuid);
+int restore_check_device(const char* uuid);
+int restore_client_new(struct idevicerestore_client_t* client);
+void restore_client_free(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);
+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 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);
+
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif
diff --git a/src/tss.c b/src/tss.c
index 4fc8b57..0d4422c 100644
--- a/src/tss.c
+++ b/src/tss.c
@@ -36,23 +36,10 @@ typedef struct {
 	char* content;
 } tss_response;
 
-plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid) {
-	// Fetch build information from BuildManifest
-	plist_t build_identities_array = plist_dict_get_item(buildmanifest, "BuildIdentities");
-	if (!build_identities_array || plist_get_node_type(build_identities_array) != PLIST_ARRAY) {
-		error("ERROR: Unable to find BuildIdentities array\n");
-		return NULL;
-	}
-
-	plist_t restore_identity_dict = plist_array_get_item(build_identities_array, 0);
-	if (!restore_identity_dict || plist_get_node_type(restore_identity_dict) != PLIST_DICT) {
-		error("ERROR: Unable to find restore identity\n");
-		return NULL;
-	}
-
+plist_t tss_create_request(plist_t build_identity, uint64_t ecid) {
 	uint64_t unique_build_size = 0;
 	char* unique_build_data = NULL;
-	plist_t unique_build_node = plist_dict_get_item(restore_identity_dict, "UniqueBuildID");
+	plist_t unique_build_node = plist_dict_get_item(build_identity, "UniqueBuildID");
 	if (!unique_build_node || plist_get_node_type(unique_build_node) != PLIST_DATA) {
 		error("ERROR: Unable to find UniqueBuildID node\n");
 		return NULL;
@@ -61,7 +48,7 @@ plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid) {
 
 	int chip_id = 0;
 	char* chip_id_string = NULL;
-	plist_t chip_id_node = plist_dict_get_item(restore_identity_dict, "ApChipID");
+	plist_t chip_id_node = plist_dict_get_item(build_identity, "ApChipID");
 	if (!chip_id_node || plist_get_node_type(chip_id_node) != PLIST_STRING) {
 		error("ERROR: Unable to find ApChipID node\n");
 		return NULL;
@@ -71,7 +58,7 @@ plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid) {
 
 	int board_id = 0;
 	char* board_id_string = NULL;
-	plist_t board_id_node = plist_dict_get_item(restore_identity_dict, "ApBoardID");
+	plist_t board_id_node = plist_dict_get_item(build_identity, "ApBoardID");
 	if (!board_id_node || plist_get_node_type(board_id_node) != PLIST_STRING) {
 		error("ERROR: Unable to find ApBoardID node\n");
 		return NULL;
@@ -81,7 +68,7 @@ plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid) {
 
 	int security_domain = 0;
 	char* security_domain_string = NULL;
-	plist_t security_domain_node = plist_dict_get_item(restore_identity_dict, "ApSecurityDomain");
+	plist_t security_domain_node = plist_dict_get_item(build_identity, "ApSecurityDomain");
 	if (!security_domain_node || plist_get_node_type(security_domain_node) != PLIST_STRING) {
 		error("ERROR: Unable to find ApSecurityDomain node\n");
 		return NULL;
@@ -112,7 +99,7 @@ plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid) {
 	free(unique_build_data);
 
 	// Add all firmware files to TSS request
-	plist_t manifest_node = plist_dict_get_item(restore_identity_dict, "Manifest");
+	plist_t manifest_node = plist_dict_get_item(build_identity, "Manifest");
 	if (!manifest_node || plist_get_node_type(manifest_node) != PLIST_DICT) {
 		error("ERROR: Unable to find restore manifest\n");
 		plist_free(tss_request);
@@ -137,13 +124,10 @@ plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid) {
 		plist_dict_insert_item(tss_request, key, tss_entry);
 		free(key);
 	}
-	plist_free(manifest_node);
 
-	int sz = 0;
-	char* xml = NULL;
-	plist_to_xml(tss_request, &xml, &sz);
-	debug("%s", xml);
-	free(xml);
+	if (idevicerestore_debug) {
+		debug_plist(tss_request);
+	}
 
 	return tss_request;
 }
@@ -198,7 +182,6 @@ plist_t tss_send_request(plist_t tss_request) {
 	curl_global_cleanup();
 
 	if (strstr(response->content, "MESSAGE=SUCCESS") == NULL) {
-		error("ERROR: Unable to get signature from this firmware\n");
 		free(response->content);
 		free(response);
 		return NULL;
@@ -219,15 +202,111 @@ plist_t tss_send_request(plist_t tss_request) {
 	free(response->content);
 	free(response);
 
-	int sz = 0;
-	char* xml = NULL;
-	plist_to_xml(tss_response, &xml, &sz);
-	debug("%s", xml);
-	free(xml);
+	if (idevicerestore_debug) {
+		debug_plist(tss_response);
+	}
 
 	return tss_response;
 }
 
-void tss_stitch_img3(img3_file* file, plist_t signature) {
+int tss_get_entry_path(plist_t tss, const char* entry, char** path) {
+	char* path_string = NULL;
+	plist_t path_node = NULL;
+	plist_t entry_node = NULL;
+
+	*path = NULL;
+
+	entry_node = plist_dict_get_item(tss, entry);
+	if (!entry_node || plist_get_node_type(entry_node) != PLIST_DICT) {
+		error("ERROR: Unable to find %s entry in TSS response\n", entry);
+		return -1;
+	}
+
+	path_node = plist_dict_get_item(entry_node, "Path");
+	if (!path_node || plist_get_node_type(path_node) != PLIST_STRING) {
+		error("ERROR: Unable to find %s path in entry\n", path_string);
+		return -1;
+	}
+	plist_get_string_val(path_node, &path_string);
+
+	*path = path_string;
+	return 0;
+}
+
+int tss_get_blob_by_path(plist_t tss, const char* path, char** blob) {
+	int i = 0;
+	uint32_t tss_size = 0;
+	uint64_t blob_size = 0;
+	char* entry_key = NULL;
+	char* blob_data = NULL;
+	char* entry_path = NULL;
+	plist_t tss_entry = NULL;
+	plist_t blob_node = NULL;
+	plist_t path_node = NULL;
+	plist_dict_iter iter = NULL;
+
+	*blob = NULL;
+
+	plist_dict_new_iter(tss, &iter);
+	tss_size = plist_dict_get_size(tss);
+	for (i = 0; i < tss_size; i++) {
+		plist_dict_next_item(tss, iter, &entry_key, &tss_entry);
+		if (entry_key == NULL)
+			break;
+
+		if (!tss_entry || plist_get_node_type(tss_entry) != PLIST_DICT) {
+			continue;
+		}
+
+		path_node = plist_dict_get_item(tss_entry, "Path");
+		if (!path_node || plist_get_node_type(path_node) != PLIST_STRING) {
+			error("ERROR: Unable to find TSS path node in entry %s\n", entry_key);
+			return -1;
+		}
+
+		plist_get_string_val(path_node, &entry_path);
+		if (strcmp(path, entry_path) == 0) {
+			blob_node = plist_dict_get_item(tss_entry, "Blob");
+			if (!blob_node || plist_get_node_type(blob_node) != PLIST_DATA) {
+				error("ERROR: Unable to find TSS blob node in entry %s\n", entry_key);
+				return -1;
+			}
+			plist_get_data_val(blob_node, &blob_data, &blob_size);
+			break;
+		}
+
+		free(entry_key);
+	}
+
+	if (blob_data == NULL || blob_size <= 0) {
+		return -1;
+	}
+
+	*blob = blob_data;
+	return 0;
+}
+
+int tss_get_blob_by_name(plist_t tss, const char* entry, char** blob) {
+	uint64_t blob_size = 0;
+	char* blob_data = NULL;
+	plist_t blob_node = NULL;
+	plist_t tss_entry = NULL;
+
+	*blob = NULL;
+
+	tss_entry = plist_dict_get_item(tss, entry);
+	if (!tss_entry || plist_get_node_type(tss_entry) != PLIST_DICT) {
+		error("ERROR: Unable to find %s entry in TSS response\n", entry);
+		return -1;
+	}
+
+	blob_node = plist_dict_get_item(tss_entry, "Blob");
+	if (!blob_node || plist_get_node_type(blob_node) != PLIST_DATA) {
+		error("ERROR: Unable to find blob in %s entry\n", entry);
+		return -1;
+	}
+	plist_get_data_val(blob_node, &blob_data, &blob_size);
 
+	*blob = blob_data;
+	return 0;
 }
diff --git a/src/tss.h b/src/tss.h
index c1ff4b4..d45c74c 100644
--- a/src/tss.h
+++ b/src/tss.h
@@ -19,15 +19,24 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  */
 
-#ifndef TSS_H
-#define TSS_H
+#ifndef IDEVICERESTORE_TSS_H
+#define IDEVICERESTORE_TSS_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 #include <plist/plist.h>
 
-#include "img3.h"
+plist_t tss_send_request(plist_t request);
+plist_t tss_create_request(plist_t build_identity, uint64_t ecid);
+int tss_get_entry_path(plist_t tss, const char* entry, char** path);
+int tss_get_blob_by_path(plist_t tss, const char* path, char** blob);
+int tss_get_blob_by_name(plist_t tss, const char* entry, char** blob);
 
-plist_t tss_create_request(plist_t buildmanifest, uint64_t ecid);
-plist_t tss_send_request(plist_t tss_request);
-void tss_stitch_img3(img3_file* file, plist_t signature);
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif