13 files changed, 1192 insertions, 415 deletions

diff --git a/heimdall/source/BridgeManager.cpp b/heimdall/source/BridgeManager.cpp
index 0b94dfb..caa3978 100644
--- a/heimdall/source/BridgeManager.cpp
+++ b/heimdall/source/BridgeManager.cpp
@@ -32,6 +32,7 @@
 #include "DumpResponse.h"
 #include "EndModemFileTransferPacket.h"
 #include "EndPhoneFileTransferPacket.h"
+#include "EndSessionPacket.h"
 #include "FileTransferPacket.h"
 #include "FlashPartFileTransferPacket.h"
 #include "FlashPartPitFilePacket.h"
@@ -41,7 +42,6 @@
 #include "PitFilePacket.h"
 #include "PitFileResponse.h"
 #include "ReceiveFilePartPacket.h"
-#include "RebootDevicePacket.h"
 #include "ResponsePacket.h"
 #include "SendFilePartPacket.h"
 #include "SendFilePartResponse.h"
@@ -54,10 +54,9 @@
 using namespace Heimdall;
 
 const DeviceIdentifier BridgeManager::supportedDevices[BridgeManager::kSupportedDeviceCount] = {
-	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidGalaxySDownloadMode)/*,
-	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidGalaxySInternational),
-	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidGalaxySNewInternational),
-	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidVibrantCanadaBell)*/
+	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidGalaxyS),
+	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidGalaxyS2),
+	DeviceIdentifier(BridgeManager::kVidSamsung, BridgeManager::kPidDroidCharge)
 };
 
 enum
@@ -428,9 +427,9 @@ bool BridgeManager::EndSession(void) const
 {
 	InterfaceManager::Print("Ending session...\n");
 
-	RebootDevicePacket *rebootDevicePacket = new RebootDevicePacket(RebootDevicePacket::kRequestEndSession);
-	bool success = SendPacket(rebootDevicePacket);
-	delete rebootDevicePacket;
+	EndSessionPacket *endSessionPacket = new EndSessionPacket(EndSessionPacket::kRequestEndSession);
+	bool success = SendPacket(endSessionPacket);
+	delete endSessionPacket;
 
 	if (!success)
 	{
@@ -438,9 +437,9 @@ bool BridgeManager::EndSession(void) const
 		return (false);
 	}
 
-	ResponsePacket *rebootDeviceResponse = new ResponsePacket(ResponsePacket::kResponseTypeRebootDevice);
-	success = ReceivePacket(rebootDeviceResponse);
-	delete rebootDeviceResponse;
+	ResponsePacket *endSessionResponse = new ResponsePacket(ResponsePacket::kResponseTypeEndSession);
+	success = ReceivePacket(endSessionResponse);
+	delete endSessionResponse;
 
 	if (!success)
 	{
@@ -539,7 +538,7 @@ bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout) const
 		}
 
 		if (verbose)
-				InterfaceManager::PrintError("\n");
+			InterfaceManager::PrintError("\n");
 	}
 
 	if (communicationDelay != 0)
@@ -1062,17 +1061,17 @@ bool BridgeManager::RebootDevice(void) const
 {
 	InterfaceManager::Print("Rebooting device...\n");
 
-	RebootDevicePacket *rebootDevicePacket = new RebootDevicePacket(RebootDevicePacket::kRequestRebootDevice);
+	EndSessionPacket *rebootDevicePacket = new EndSessionPacket(EndSessionPacket::kRequestRebootDevice);
 	bool success = SendPacket(rebootDevicePacket);
 	delete rebootDevicePacket;
 
 	if (!success)
 	{
-		InterfaceManager::PrintError("Failed to send end session packet!\n");
+		InterfaceManager::PrintError("Failed to send reboot device packet!\n");
 		return (false);
 	}
 
-	ResponsePacket *rebootDeviceResponse = new ResponsePacket(ResponsePacket::kResponseTypeRebootDevice);
+	ResponsePacket *rebootDeviceResponse = new ResponsePacket(ResponsePacket::kResponseTypeEndSession);
 	success = ReceivePacket(rebootDeviceResponse);
 	delete rebootDeviceResponse;
 
diff --git a/heimdall/source/BridgeManager.h b/heimdall/source/BridgeManager.h
index 40a7ceb..06e4d61 100644
--- a/heimdall/source/BridgeManager.h
+++ b/heimdall/source/BridgeManager.h
@@ -53,7 +53,7 @@ namespace Heimdall
 

 			enum

 			{

-				kSupportedDeviceCount		= 1,

+				kSupportedDeviceCount		= 3,

 

 				kCommunicationDelayDefault	= 0,

 				kDumpBufferSize				= 4096

@@ -66,10 +66,9 @@ namespace Heimdall
 

 			enum

 			{

-				kPidGalaxySDownloadMode		= 0x6601/*,

-				kPidGalaxySInternational	= 0x681C,

-				kPidGalaxySNewInternational	= 0x681D,

-				kPidVibrantCanadaBell		= 0x6877*/

+				kPidGalaxyS		    = 0x6601,

+				kPidGalaxyS2        = 0x685D,

+				kPidDroidCharge     = 0x68C3

 			};

 

 		private:

@@ -105,8 +104,8 @@ namespace Heimdall
 

 			bool SendPacket(OutboundPacket *packet, int timeout = 3000) const;

 			bool ReceivePacket(InboundPacket *packet, int timeout = 3000) const;

-
-			bool SendPitFile(FILE *file) const;
+

+			bool SendPitFile(FILE *file) const;

 			int ReceivePitFile(unsigned char **pitBuffer) const;

 

 			bool SendFile(FILE *file, int destination, int fileIdentifier = -1) const;

diff --git a/heimdall/source/ControlPacket.h b/heimdall/source/ControlPacket.h
index 2379073..019513e 100644
--- a/heimdall/source/ControlPacket.h
+++ b/heimdall/source/ControlPacket.h
@@ -32,10 +32,10 @@ namespace Heimdall
 

 			enum

 			{

-				kControlTypeDeviceInfo		= 0x64,

-				kControlTypePitFile			= 0x65,

-				kControlTypeFileTransfer	= 0x66,

-				kControlTypeRebootDevice	= 0x67

+				kControlTypeDeviceInfo      = 0x64,

+				kControlTypePitFile         = 0x65,

+				kControlTypeFileTransfer    = 0x66,

+				kControlTypeEndSession      = 0x67

 			};

 

 		protected:

diff --git a/heimdall/source/EndPhoneFileTransferPacket.h b/heimdall/source/EndPhoneFileTransferPacket.h
index 24650a2..f875002 100644
--- a/heimdall/source/EndPhoneFileTransferPacket.h
+++ b/heimdall/source/EndPhoneFileTransferPacket.h
@@ -30,7 +30,7 @@ namespace Heimdall
 	{

 		public:

 

-			enum

+			/*enum

 			{

 				kFilePrimaryBootloader			= 0x00,

 				kFilePit						= 0x01, // New 1.1 - Don't flash the pit this way!

@@ -38,6 +38,9 @@ namespace Heimdall
 				kFileSecondaryBootloaderBackup	= 0x04,	// New 1.1

 				kFileKernel						= 0x06,

 				kFileRecovery					= 0x07,	// New 1.1

+

+				kFileTabletModem				= 0x08, // New 1.2

+

 				kFileEfs						= 0x14, // New 1.1

 				kFileParamLfs					= 0x15,

 				kFileFactoryFilesystem			= 0x16,

@@ -45,7 +48,7 @@ namespace Heimdall
 				kFileCache						= 0x18,

 

 				kFileModem						= 0x0B	// New 1.1 - Kies flashes the modem this way rather than using the EndModemFileTransferPacket.

-			};

+			};*/

 

 		private:

 

diff --git a/heimdall/source/RebootDevicePacket.h b/heimdall/source/EndSessionPacket.h
index f40a227..15c7162 100644
--- a/heimdall/source/RebootDevicePacket.h
+++ b/heimdall/source/EndSessionPacket.h
@@ -1,64 +1,64 @@
-/* Copyright (c) 2010 Benjamin Dobell, Glass Echidna

- 

- Permission is hereby granted, free of charge, to any person obtaining a copy

- of this software and associated documentation files (the "Software"), to deal

- in the Software without restriction, including without limitation the rights

- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

- copies of the Software, and to permit persons to whom the Software is

- furnished to do so, subject to the following conditions:

- 

- The above copyright notice and this permission notice shall be included in

- all copies or substantial portions of the Software.

- 

- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

- THE SOFTWARE.*/

-

-#ifndef REBOOTDEVICEPACKET_H

-#define REBOOTDEVICEPACKET_H

-

-// Heimdall

-#include "ControlPacket.h"

-

-namespace Heimdall

-{

-	class RebootDevicePacket : public ControlPacket

-	{

-		public:

-

-			enum

-			{

-				kRequestEndSession = 0,

-				kRequestRebootDevice = 1

-			};

-

-		private:

-

-			unsigned int request;

-

-		public:

-

-			RebootDevicePacket(unsigned int request) : ControlPacket(ControlPacket::kControlTypeRebootDevice)

-			{

-				this->request = request;

-			}

-

-			unsigned int GetRequest(void) const

-			{

-				return (request);

-			}

-

-			void Pack(void)

-			{

-				ControlPacket::Pack();

-

-				PackInteger(ControlPacket::kDataSize, request);

-			}

-	};

-}

-

-#endif

+/* Copyright (c) 2010 Benjamin Dobell, Glass Echidna
+ 
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+ 
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+ 
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE.*/
+
+#ifndef REBOOTDEVICEPACKET_H
+#define REBOOTDEVICEPACKET_H
+
+// Heimdall
+#include "ControlPacket.h"
+
+namespace Heimdall
+{
+	class EndSessionPacket : public ControlPacket
+	{
+		public:
+
+			enum
+			{
+				kRequestEndSession = 0,
+				kRequestRebootDevice = 1
+			};
+
+		private:
+
+			unsigned int request;
+
+		public:
+
+			EndSessionPacket(unsigned int request) : ControlPacket(ControlPacket::kControlTypeEndSession)
+			{
+				this->request = request;
+			}
+
+			unsigned int GetRequest(void) const
+			{
+				return (request);
+			}
+
+			void Pack(void)
+			{
+				ControlPacket::Pack();
+
+				PackInteger(ControlPacket::kDataSize, request);
+			}
+	};
+}
+
+#endif
diff --git a/heimdall/source/Heimdall.h b/heimdall/source/Heimdall.h
index 672cf14..174e7d7 100644
--- a/heimdall/source/Heimdall.h
+++ b/heimdall/source/Heimdall.h
@@ -21,19 +21,19 @@
 #ifndef HEIMDALL_H

 #define HEIMDALL_H

 

-#ifdef OS_WINDOWS
-#include <Windows.h>
-#else
-
-#include "../config.h"
-
-#if defined(OS_DARWIN) || defined(OS_LINUX)
-#include <unistd.h>
-#define Sleep(t) usleep(1000*t)
-#else
-#error operating system not supported
-#endif
-
+#ifdef OS_WINDOWS

+#include <Windows.h>

+#else

+

+#include "../config.h"

+

+#if defined(OS_DARWIN) || defined(OS_LINUX)

+#include <unistd.h>

+#define Sleep(t) usleep(1000*t)

+#else

+#error operating system not supported

+#endif

+

 #endif

 

 #ifndef nullptr

diff --git a/heimdall/source/InboundPacket.h b/heimdall/source/InboundPacket.h
index 6ff3087..3285c48 100644
--- a/heimdall/source/InboundPacket.h
+++ b/heimdall/source/InboundPacket.h
@@ -35,14 +35,14 @@ namespace Heimdall
 

 		protected:

 

-			int UnpackInteger(int offset)

+			unsigned int UnpackInteger(unsigned int offset) const

 			{

 #ifdef WORDS_BIGENDIAN

-				int value = (data[offset] << 24) | (data[offset + 1] << 16) |

+				unsigned int value = (data[offset] << 24) | (data[offset + 1] << 16) |

 					(data[offset + 2] << 8) | data[offset + 3];

 #else

 				// Flip endianness

-				int value = data[offset] | (data[offset + 1] << 8) |

+				unsigned int value = data[offset] | (data[offset + 1] << 8) |

 					(data[offset + 2] << 16) | (data[offset + 3] << 24);

 #endif

 				return (value);

diff --git a/heimdall/source/InterfaceManager.cpp b/heimdall/source/InterfaceManager.cpp
index 23791bc..aaf6d0b 100644
--- a/heimdall/source/InterfaceManager.cpp
+++ b/heimdall/source/InterfaceManager.cpp
@@ -29,18 +29,20 @@
 using namespace std;

 using namespace Heimdall;

 

-string InterfaceManager::actionNames[kActionCount] = { "flash", "close-pc-screen", "dump", "help" };

+string InterfaceManager::actionNames[kActionCount] = { "flash", "close-pc-screen", "dump", "print-pit", "help" };

 

 string InterfaceManager::flashArgumentNames[kFlashArgCount * 2] = {

 	// --- Long Names ---

 	"-repartition",

 

 	"-pit", "-factoryfs", "-cache", "-dbdata", "-primary-boot",	"-secondary-boot", "-secondary-boot-backup", "-param", "-kernel", "-recovery", "-efs", "-modem",

+	"-normal-boot", "-system", "-user-data", "-fota", "-hidden", "-movinand", "-data", "-ums", "-emmc", "-%d",

 

 	// --- Short Names ---

 	"r",

 

-	"pit",  "fs",         "cache",  "db",      "boot",           "sbl",            "sbl2",                   "param",  "z",       "rec",       "efs",  "m"

+	"pit",  "fs",         "cache",  "db",      "boot",           "sbl",            "sbl2",                   "param",  "z",       "rec",       "efs",  "m",

+	"norm",         "sys",     "udata",      "fota",  "hide",    "nand",      "data",  "ums",  "emmc",  "%d"

 };

 

 string InterfaceManager::dumpArgumentNames[kDumpArgCount * 2] = {

@@ -53,12 +55,12 @@ string InterfaceManager::dumpArgumentNames[kDumpArgCount * 2] = {
 

 string InterfaceManager::commonArgumentNames[kCommonArgCount * 2] = {

 	// --- Long Names ---

-	"-verbose",

+	"-verbose", "-no-reboot"

 

 	"-delay",

 

 	// --- Short Names ---

-	"v",

+	"v",        "nobt"

 

 	"d"

 };

@@ -74,6 +76,9 @@ string *InterfaceManager::actionArgumentNames[kActionCount + 1] = {
 	// kActionDump

 	dumpArgumentNames,

 

+	// kActionPrintPit

+	nullptr,

+

 	// kActionHelp

 	nullptr,

 

@@ -82,28 +87,36 @@ string *InterfaceManager::actionArgumentNames[kActionCount + 1] = {
 };

 

 int InterfaceManager::actionArgumentCounts[kActionCount + 1] = {

-	kFlashArgCount, 0, kDumpArgCount, 0, kCommonArgCount

+	kFlashArgCount, 0, kDumpArgCount, 0, 0, kCommonArgCount

 };

 

 int InterfaceManager::actionValuelessArgumentCounts[kActionCount + 1] = {

-	kFlashArgPit, 0, kDumpArgChipType, 0, kCommonArgDelay

+	kFlashArgPit, 0, kDumpArgChipType, 0, 0, kCommonArgDelay

 };

 

-const char *InterfaceManager::usage = "Usage: heimdall <action> <arguments> [--verbose] [--delay <ms>]\n\

+const char *InterfaceManager::usage = "Usage: heimdall <action> <arguments> [--verbose] [--no-reboot] [--delay <ms>]\n\

 \n\

 action: flash\n\

 arguments:\n\

-    --repartition --pit <filename> --factoryfs <filename>\n\

-    --cache <filename> --dbdata <filename> --primary-boot <filename>\n\

-    --secondary-boot <filename> --param <filename> --kernel <filename>\n\

-    --modem <filename>\n\

+    --repartition --pit <filename> [--factoryfs <filename>]\n\

+    [--cache <filename>] [--dbdata <filename>] [--primary-boot <filename>]\n\

+    [--secondary-boot <filename>] [--param <filename>] [--kernel <filename>]\n\

+	[--modem <filename>] [--normal-boot <filename>] [--system <filename>]\n\

+	[--user-data <filename>] [--fota <filename>] [--hidden <filename>]\n\

+	[--movinand <filename>] [--data <filename>] [--ums <filename>]\n\

+	[--emmc <filename>] [--<partition identifier> <filename>]\n\

   or:\n\

-    [--pit <filename>] [--factoryfs <filename>] [--cache <filename>]\n\

-    [--dbdata <filename>] [--primary-boot <filename>]\n\

-    [--secondary-boot <filename>] [--secondary-boot-backup <filename>]\n\

-    [--param <filename>] [--kernel <filename>] [--recovery <filename>]\n\

-    [--efs <filename>] [--modem <filename>]\n\

+    [--factoryfs <filename>] [--cache <filename>] [--dbdata <filename>]\n\

+    [--primary-boot <filename>] [--secondary-boot <filename>]\n\

+    [--secondary-boot-backup <filename>] [--param <filename>]\n\

+    [--kernel <filename>] [--recovery <filename>] [--efs <filename>]\n\

+    [--modem <filename>] [--normal-boot <filename>] [--system <filename>]\n\

+	[--user-data <filename>] [--fota <filename>] [--hidden <filename>]\n\

+	[--movinand <filename>] [--data <filename>] [--ums <filename>]\n\

+	[--emmc <filename>] [--<partition identifier> <filename>]\n\

 description: Flashes firmware files to your phone.\n\

+WARNING: If you're repartitioning it's strongly recommended you specify\n\

+         all files at your disposal, including bootloaders.\n\

 \n\

 action: close-pc-screen\n\

 description: Attempts to get rid off the \"connect phone to PC\" screen.\n\

@@ -114,6 +127,10 @@ description: Attempts to dump data from the phone corresponding to the\n\
 	specified chip type and chip ID.\n\

 NOTE: Galaxy S phones don't appear to properly support this functionality.\n\

 \n\

+action: print-pit\n\

+description: Dumps the PIT file from the connected device and prints it in\n\

+    a human readable format.\n\

+\n\

 action: help\n\

 description: Display this dialogue.\n";

 

@@ -200,6 +217,25 @@ bool InterfaceManager::GetArguments(int argc, char **argv, map<string, string>&
 		// Check if the argument is a valid regular argument

 		for (int i = actionValuelessArgumentCount; i < actionArgumentCount; i++)

 		{

+			// Support for --<integer> and -<integer> parameters.

+			if (argumentName.length() > 1 && argumentNames[i].compare("-%d") == 0)

+			{

+				if (atoi(argumentName.substr(1).c_str()) > 0 || argumentName.compare("-0") == 0)

+				{

+					valid = true;

+					break;

+				}

+			}

+			else if (argumentNames[i].compare("%d") == 0)

+			{

+				if (atoi(argumentName.c_str()) > 0 || argumentName.compare("0") == 0)

+				{

+					argumentName = "-" + argumentName;

+					valid = true;

+					break;

+				}

+			}

+

 			if (argumentName == argumentNames[i] || argumentName == argumentNames[actionArgumentCount + i])

 			{

 				argumentName = argumentNames[i];

@@ -213,6 +249,25 @@ bool InterfaceManager::GetArguments(int argc, char **argv, map<string, string>&
 			// Check if it's a common regular argument

 			for (int i = commonValuelessArgumentCount; i < commonArgumentCount; i++)

 			{

+				// Support for --<integer> and -<integer> parameters.

+				if (argumentName.length() > 1 && argumentNames[i].compare("-%d"))

+				{

+					if (atoi(argumentName.substr(1).c_str()) > 0 || argumentName.compare("-0") == 0)

+					{

+						valid = true;

+						break;

+					}

+				}

+				else if (argumentNames[i].compare("%d"))

+				{

+					if (atoi(argumentName.c_str()) > 0 || argumentName.compare("0") == 0)

+					{

+						argumentName = "-" + argumentName;

+						valid = true;

+						break;

+					}

+				}

+

 				if (argumentName == commonArgumentNames[i] || argumentName == commonArgumentNames[commonArgumentCount + i])

 				{

 					argumentName = commonArgumentNames[i];

diff --git a/heimdall/source/InterfaceManager.h b/heimdall/source/InterfaceManager.h
index 734c086..1f5ba94 100644
--- a/heimdall/source/InterfaceManager.h
+++ b/heimdall/source/InterfaceManager.h
@@ -38,6 +38,7 @@ namespace Heimdall
 				kActionFlash = 0,
 				kActionClosePcScreen,
 				kActionDump,
+				kActionPrintPit,
 				kActionHelp,
 				kActionCount
 			};
@@ -51,7 +52,7 @@ namespace Heimdall
 				kFlashArgPit,
 				kFlashArgFactoryFs,
 				kFlashArgCache,
-				kFlashArgData,
+				kFlashArgDatabaseData,
 				kFlashArgPrimaryBootloader,
 				kFlashArgSecondaryBootloader,
 				kFlashArgSecondaryBootloaderBackup,
@@ -61,6 +62,18 @@ namespace Heimdall
 				kFlashArgEfs,
 				kFlashArgModem,
 
+				kFlashArgNormalBoot,
+				kFlashArgSystem,
+				kFlashArgUserData,
+				kFlashArgFota,
+				kFlashArgHidden,
+				kFlashArgMovinand,
+				kFlashArgData,
+				kFlashArgUms,
+				kFlashArgEmmc,
+
+				kFlashArgPartitionIndex,
+
 				kFlashArgCount
 			};
 
@@ -78,6 +91,7 @@ namespace Heimdall
 			{
 				// Valueless arguments
 				kCommonArgVerbose = 0,
+				kCommonArgNoReboot,
 
 				// Regular arguments
 				kCommonArgDelay,
diff --git a/heimdall/source/PitData.cpp b/heimdall/source/PitData.cpp
new file mode 100755
index 0000000..10e88b7
--- /dev/null
+++ b/heimdall/source/PitData.cpp
@@ -0,0 +1,284 @@
+/* Copyright (c) 2010 Benjamin Dobell, Glass Echidna
+ 
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+ 
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+ 
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE.*/
+
+// Heimdall
+#include "Heimdall.h"
+#include "InterfaceManager.h"
+#include "PitData.h"
+
+using namespace Heimdall;
+
+PitEntry::PitEntry()
+{
+	unused = false;
+	partitionType = 0;
+	partitionIdentifier = 0;
+	partitionFlags = 0;
+	unknown2 = 0;
+	partitionBlockSize = 0;
+	partitionBlockCount = 0;
+	unknown3 = 0;
+	unknown4 = 0;
+
+	memset(partitionName, 0, 32);
+	memset(filename, 0, 64);
+}
+
+PitEntry::~PitEntry()
+{
+}
+
+void PitEntry::Print(void) const
+{
+	InterfaceManager::Print("Unused: %s\n", (unused) ? "Yes" : "No");
+
+	const char *partitionTypeText = "Unknown";
+
+	if (partitionType == PitEntry::kPartitionTypeRfs)
+		partitionTypeText = "RFS";
+	else if (partitionType == PitEntry::kPartitionTypeExt4)
+		partitionTypeText = "EXT4";
+
+	InterfaceManager::Print("Partition Type: %d (%s)\n", partitionType, partitionTypeText);
+
+	InterfaceManager::Print("Partition Identifier: %d\n", partitionIdentifier);
+
+	InterfaceManager::Print("Partition Flags: %d (", partitionFlags);
+
+	if (partitionFlags & PitEntry::kPartitionFlagWrite)
+		InterfaceManager::Print("R/W");
+	else
+		InterfaceManager::Print("R");
+
+	InterfaceManager::Print(")\n");
+
+	InterfaceManager::Print("Unknown 2: %d\n", unknown2);
+
+	InterfaceManager::Print("Partition Block Size: %d\n", partitionBlockSize);
+	InterfaceManager::Print("Partition Block Count: %d\n", partitionBlockCount);
+
+	InterfaceManager::Print("Unknown 3: %d\n", unknown3);
+	InterfaceManager::Print("Unknown 4: %d\n", unknown4);
+
+	InterfaceManager::Print("Partition Name: %s\n", partitionName);
+	InterfaceManager::Print("Filename: %s\n", filename);
+}
+
+
+
+PitData::PitData()
+{
+	entryCount = 0;
+
+	unknown1 = 0;
+	unknown2 = 0;
+
+	unknown3 = 0;
+	unknown4 = 0;
+
+	unknown5 = 0;
+	unknown6 = 0;
+
+	unknown7 = 0;
+	unknown8 = 0;
+}
+
+PitData::~PitData()
+{
+	for (unsigned int i = 0; i < entries.size(); i++)
+		delete entries[i];
+}
+
+bool PitData::Unpack(const unsigned char *data)
+{
+	if (PitData::UnpackInteger(data, 0) != PitData::kFileIdentifier)
+		return (false);
+
+	// Remove existing entries
+	for (unsigned int i = 0; i < entries.size(); i++)
+		delete entries[i];
+
+	entryCount = PitData::UnpackInteger(data, 4);
+
+	entries.resize(entryCount);
+
+	unknown1 = PitData::UnpackInteger(data, 8);
+	unknown2 = PitData::UnpackInteger(data, 12);
+
+	unknown3 = PitData::UnpackShort(data, 16);
+	unknown4 = PitData::UnpackShort(data, 18);
+
+	unknown5 = PitData::UnpackShort(data, 20);
+	unknown6 = PitData::UnpackShort(data, 22);
+
+	unknown7 = PitData::UnpackShort(data, 24);
+	unknown8 = PitData::UnpackShort(data, 26);
+
+	unsigned int integerValue;
+	unsigned int entryOffset;
+
+	for (unsigned int i = 0; i < entryCount; i++)
+	{
+		entryOffset = PitData::kHeaderDataSize + i * PitEntry::kDataSize;
+
+		entries[i] = new PitEntry();
+
+		integerValue = PitData::UnpackInteger(data, entryOffset);
+		entries[i]->SetUnused((integerValue != 0) ? true : false);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 4);
+		entries[i]->SetPartitionType(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 8);
+		entries[i]->SetPartitionIdentifier(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 12);
+		entries[i]->SetPartitionFlags(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 16);
+		entries[i]->SetUnknown2(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 20);
+		entries[i]->SetPartitionBlockSize(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 24);
+		entries[i]->SetPartitionBlockCount(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 28);
+		entries[i]->SetUnknown3(integerValue);
+
+		integerValue = PitData::UnpackInteger(data, entryOffset + 32);
+		entries[i]->SetUnknown4(integerValue);
+
+		entries[i]->SetPartitionName((const char *)data + entryOffset + 36);
+		entries[i]->SetFilename((const char *)data + entryOffset + 36 + PitEntry::kPartitionNameMaxLength);
+	}
+
+	return (true);
+}
+
+void PitData::Pack(unsigned char *data) const
+{
+	PitData::PackInteger(data, 0, PitData::kFileIdentifier);
+
+	PitData::PackInteger(data, 4, entryCount);
+
+	PitData::PackInteger(data, 8, unknown1);
+	PitData::PackInteger(data, 12, unknown2);
+
+	PitData::PackShort(data, 16, unknown3);
+	PitData::PackShort(data, 18, unknown4);
+
+	PitData::PackShort(data, 20, unknown5);
+	PitData::PackShort(data, 22, unknown6);
+
+	PitData::PackShort(data, 24, unknown7);
+	PitData::PackShort(data, 26, unknown8);
+
+	int entryOffset;
+
+	for (unsigned int i = 0; i < entryCount; i++)
+	{
+		entryOffset = PitData::kHeaderDataSize + i * PitEntry::kDataSize;
+
+		PitData::PackInteger(data, entryOffset, (entries[i]->GetUnused()) ? 1 : 0);
+
+		PitData::PackInteger(data, entryOffset + 4, entries[i]->GetPartitionType());
+		PitData::PackInteger(data, entryOffset + 8, entries[i]->GetPartitionIdentifier());
+		PitData::PackInteger(data, entryOffset + 12, entries[i]->GetPartitionFlags());
+
+		PitData::PackInteger(data, entryOffset + 16, entries[i]->GetUnknown2());
+
+		PitData::PackInteger(data, entryOffset + 20, entries[i]->GetPartitionBlockSize());
+		PitData::PackInteger(data, entryOffset + 24, entries[i]->GetPartitionBlockCount());
+
+		PitData::PackInteger(data, entryOffset + 28, entries[i]->GetUnknown3());
+		PitData::PackInteger(data, entryOffset + 32, entries[i]->GetUnknown4());
+
+		memcpy(data + entryOffset + 36, entries[i]->GetPartitionName(), PitEntry::kPartitionNameMaxLength);
+		memcpy(data + entryOffset + 36 + PitEntry::kPartitionNameMaxLength, entries[i]->GetPartitionName(), PitEntry::kFilenameMaxLength);
+	}
+}
+
+PitEntry *PitData::FindEntry(const char *partitionName)
+{
+	for (unsigned int i = 0; i < entries.size(); i++)
+	{
+		if (strcmp(entries[i]->GetPartitionName(), partitionName) == 0)
+			return (entries[i]);
+	}
+
+	return (nullptr);
+}
+
+const PitEntry *PitData::FindEntry(const char *partitionName) const
+{
+	for (unsigned int i = 0; i < entries.size(); i++)
+	{
+		if (strcmp(entries[i]->GetPartitionName(), partitionName) == 0)
+			return (entries[i]);
+	}
+
+	return (nullptr);
+}
+
+PitEntry *PitData::FindEntry(unsigned int partitionIdentifier)
+{
+	for (unsigned int i = 0; i < entries.size(); i++)
+	{
+		if (entries[i]->GetPartitionIdentifier() == partitionIdentifier)
+			return (entries[i]);
+	}
+
+	return (nullptr);
+}
+
+const PitEntry *PitData::FindEntry(unsigned int partitionIdentifier) const
+{
+	for (unsigned int i = 0; i < entries.size(); i++)
+	{
+		if (entries[i]->GetPartitionIdentifier() == partitionIdentifier)
+			return (entries[i]);
+	}
+
+	return (nullptr);
+}
+
+void PitData::Print(void) const
+{
+	InterfaceManager::Print("Entry Count: %d\n", entryCount);
+
+	InterfaceManager::Print("Unknown 1: %d\n", unknown1);
+	InterfaceManager::Print("Unknown 2: %d\n", unknown2);
+	InterfaceManager::Print("Unknown 3: %d\n", unknown3);
+	InterfaceManager::Print("Unknown 4: %d\n", unknown4);
+	InterfaceManager::Print("Unknown 5: %d\n", unknown5);
+	InterfaceManager::Print("Unknown 6: %d\n", unknown6);
+	InterfaceManager::Print("Unknown 7: %d\n", unknown7);
+	InterfaceManager::Print("Unknown 8: %d\n", unknown8);
+
+	for (unsigned int i = 0; i < entryCount; i++)
+	{
+		InterfaceManager::Print("\n\n--- Entry #%d ---\n", i);
+		entries[i]->Print();
+	}
+
+	InterfaceManager::Print("\n");
+}
diff --git a/heimdall/source/PitData.h b/heimdall/source/PitData.h
new file mode 100755
index 0000000..b6a0d2f
--- /dev/null
+++ b/heimdall/source/PitData.h
@@ -0,0 +1,345 @@
+/* Copyright (c) 2010 Benjamin Dobell, Glass Echidna
+ 
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+ 
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+ 
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE.*/
+
+#ifndef PITDATA_H
+#define PITDATA_H
+
+// C Standard Library
+#include <string.h>
+#include <vector>
+
+namespace Heimdall
+{
+	class PitEntry
+	{
+		public:
+
+			enum
+			{
+				kDataSize = 132,
+				kPartitionNameMaxLength = 32,
+				kFilenameMaxLength = 64
+			};
+
+			enum
+			{
+				kPartitionTypeRfs       = 0,
+				kPartitionTypeBlank     = 1, // ?
+				kPartitionTypeExt4      = 2
+			};
+
+			enum
+			{
+				kPartitionFlagWrite     = 1 << 1
+			};
+
+		private:
+
+			bool unused;
+
+			unsigned int partitionType;
+			unsigned int partitionIdentifier;
+			unsigned int partitionFlags;
+
+			unsigned int unknown2;
+
+			unsigned int partitionBlockSize;
+			unsigned int partitionBlockCount;
+
+			unsigned int unknown3;
+			unsigned int unknown4;
+
+			char partitionName[kPartitionNameMaxLength];
+			char filename[kFilenameMaxLength];
+
+		public:
+
+			PitEntry();
+			~PitEntry();
+
+			void Print(void) const;
+
+			bool GetUnused(void) const
+			{
+				return unused;
+			}
+
+			void SetUnused(bool unused)
+			{
+				this->unused = unused;
+			}
+
+			unsigned int GetPartitionType(void) const
+			{
+				return partitionType;
+			}
+
+			void SetPartitionType(unsigned int partitionType)
+			{
+				this->partitionType = partitionType;
+			}
+
+			unsigned int GetPartitionIdentifier(void) const
+			{
+				return partitionIdentifier;
+			}
+
+			void SetPartitionIdentifier(unsigned int partitionIdentifier)
+			{
+				this->partitionIdentifier = partitionIdentifier;
+			}
+
+			unsigned int GetPartitionFlags(void) const
+			{
+				return partitionFlags;
+			}
+
+			void SetPartitionFlags(unsigned int partitionFlags)
+			{
+				this->partitionFlags = partitionFlags;
+			}
+
+			unsigned int GetUnknown2(void) const
+			{
+				return unknown2;
+			}
+
+			void SetUnknown2(unsigned int unknown2)
+			{
+				this->unknown2 = unknown2;
+			}
+
+			unsigned int GetPartitionBlockSize(void) const
+			{
+				return partitionBlockSize;
+			}
+
+			void SetPartitionBlockSize(unsigned int partitionBlockSize)
+			{
+				this->partitionBlockSize = partitionBlockSize;
+			}
+
+			unsigned int GetPartitionBlockCount(void) const
+			{
+				return partitionBlockCount;
+			}
+
+			void SetPartitionBlockCount(unsigned int partitionBlockCount)
+			{
+				this->partitionBlockCount = partitionBlockCount;
+			}
+
+			unsigned int GetUnknown3(void) const
+			{
+				return unknown3;
+			}
+
+			void SetUnknown3(unsigned int unknown3)
+			{
+				this->unknown3 = unknown3;
+			}
+
+			unsigned int GetUnknown4(void) const
+			{
+				return unknown4;
+			}
+
+			void SetUnknown4(unsigned int unknown4)
+			{
+				this->unknown4 = unknown4;
+			}
+
+			const char *GetPartitionName(void) const
+			{
+				return partitionName;
+			}
+
+			void SetPartitionName(const char *partitionName)
+			{
+				// This isn't strictly necessary but ensures no junk is left in our PIT file.
+				memset(this->partitionName, 0, 64);
+
+				if (strlen(partitionName) < 64)
+					strcpy(this->partitionName, partitionName);
+				else
+					memcpy(this->partitionName, partitionName, 63);
+			}
+
+			const char *GetFilename(void) const
+			{
+				return filename;
+			}
+
+			void SetFilename(const char *filename)
+			{
+				// This isn't strictly necessary but ensures no junk is left in our PIT file.
+				memset(this->filename, 0, 32);
+
+				if (strlen(partitionName) < 32)
+					strcpy(this->filename, filename);
+				else
+					memcpy(this->filename, filename, 31);
+			}
+	};
+
+	class PitData
+	{
+		public:
+
+			enum
+			{
+				kFileIdentifier = 0x12349876,
+				kHeaderDataSize = 28
+			};
+
+		private:
+
+			unsigned int entryCount; // 0x04
+			unsigned int unknown1;   // 0x08
+			unsigned int unknown2;   // 0x0C
+
+			unsigned short unknown3; // 0x10 (7508 = I9000, 7703 = I9100 & P1000)?
+			unsigned short unknown4; // 0x12 (Always 65, probably flags of some sort)
+
+			unsigned short unknown5; // 0x14
+			unsigned short unknown6; // 0x16
+
+			unsigned short unknown7; // 0x18
+			unsigned short unknown8; // 0x1A
+
+			// Entries start at 0x1C
+			std::vector<PitEntry *> entries;
+
+			static int UnpackInteger(const unsigned char *data, unsigned int offset)
+			{
+#ifdef WORDS_BIGENDIAN
+				int value = (data[offset] << 24) | (data[offset + 1] << 16) |
+					(data[offset + 2] << 8) | data[offset + 3];
+#else
+				// Flip endianness
+				int value = data[offset] | (data[offset + 1] << 8) |
+					(data[offset + 2] << 16) | (data[offset + 3] << 24);
+#endif
+				return (value);
+			}
+
+			static int UnpackShort(const unsigned char *data, unsigned int offset)
+			{
+#ifdef WORDS_BIGENDIAN
+				short value = (data[offset] << 8) | data[offset + 1];
+#else
+				// Flip endianness
+				short value = data[offset] | (data[offset + 1] << 8);
+#endif
+				return (value);
+			}
+
+			static void PackInteger(unsigned char *data, unsigned int offset, unsigned int value)
+			{
+#ifdef WORDS_BIGENDIAN
+				data[offset] = (value & 0xFF000000) >> 24;
+				data[offset + 1] = (value & 0x00FF0000) >> 16;
+				data[offset + 2] = (value & 0x0000FF00) >> 8;
+				data[offset + 3] = value & 0x000000FF;
+#else
+				// Flip endianness
+				data[offset] = value & 0x000000FF;
+				data[offset + 1] = (value & 0x0000FF00) >> 8;
+				data[offset + 2] = (value & 0x00FF0000) >> 16;
+				data[offset + 3] = (value & 0xFF000000) >> 24;
+#endif
+			}
+
+			static void PackShort(unsigned char *data, unsigned int offset, unsigned short value)
+			{
+#ifdef WORDS_BIGENDIAN
+				data[offset] = (value & 0xFF00) >> 8;
+				data[offset + 1] = value & 0x00FF;
+#else
+				// Flip endianness
+				data[offset] = value & 0x00FF;
+				data[offset + 1] = (value & 0xFF00) >> 8;
+#endif
+			}
+
+		public:
+
+			PitData();
+			~PitData();
+
+			bool Unpack(const unsigned char *data);
+			void Pack(unsigned char *data) const;
+
+			PitEntry *FindEntry(const char *partitionName);
+			const PitEntry *FindEntry(const char *partitionName) const;
+
+			PitEntry *FindEntry(unsigned int partitionIdentifier);
+			const PitEntry *FindEntry(unsigned int partitionIdentifier) const;
+
+			void Print(void) const;
+
+			unsigned int GetEntryCount(void) const
+			{
+				return entryCount;
+			}
+
+			unsigned int GetUnknown1(void) const
+			{
+				return unknown1;
+			}
+
+			unsigned int GetUnknown2(void) const
+			{
+				return unknown2;
+			}
+
+			unsigned short GetUnknown3(void) const
+			{
+				return unknown3;
+			}
+
+			unsigned short GetUnknown4(void) const
+			{
+				return unknown4;
+			}
+
+			unsigned short GetUnknown5(void) const
+			{
+				return unknown5;
+			}
+
+			unsigned short GetUnknown6(void) const
+			{
+				return unknown6;
+			}
+
+			unsigned short GetUnknown7(void) const
+			{
+				return unknown7;
+			}
+
+			unsigned short GetUnknown8(void) const
+			{
+				return unknown8;
+			}
+	};
+}
+
+#endif
diff --git a/heimdall/source/ResponsePacket.h b/heimdall/source/ResponsePacket.h
index ce799b5..8cafc95 100644
--- a/heimdall/source/ResponsePacket.h
+++ b/heimdall/source/ResponsePacket.h
@@ -32,11 +32,11 @@ namespace Heimdall
 

 			enum

 			{

-				kResponseTypeSendFilePart	= 0x00,

-				kResponseTypeDeviceInfo		= 0x64,

-				kResponseTypePitFile		= 0x65,

-				kResponseTypeFileTransfer	= 0x66,

-				kResponseTypeRebootDevice	= 0x67

+				kResponseTypeSendFilePart   = 0x00,

+				kResponseTypeDeviceInfo     = 0x64,

+				kResponseTypePitFile        = 0x65,

+				kResponseTypeFileTransfer   = 0x66,

+				kResponseTypeEndSession     = 0x67

 			};

 

 		private:

diff --git a/heimdall/source/main.cpp b/heimdall/source/main.cpp
index 7fa00be..81e1ea7 100644
--- a/heimdall/source/main.cpp
+++ b/heimdall/source/main.cpp
@@ -31,220 +31,288 @@
 #include "EndModemFileTransferPacket.h"
 #include "EndPhoneFileTransferPacket.h"
 #include "InterfaceManager.h"
+#include "PitData.h"
 
 using namespace std;
 using namespace Heimdall;
 
+// Known partitions
 enum
 {
-	kFilePit = 0,
-	kFileFactoryFs,
-	kFileCache,
-	kFileData,
-	kFilePrimaryBootloader,
-	kFileSecondaryBootloader,
-	kFileSecondaryBootloaderBackup,
-	kFileParam,
-	kFileKernel,
-	kFileRecovery,
-	kFileEfs,
-	kFileModem,
-	kFileCount
+	kKnownPartitionPit = 0,
+	kKnownPartitionFactoryFs,
+	kKnownPartitionCache,
+	kKnownPartitionDatabaseData,
+	kKnownPartitionPrimaryBootloader,
+	kKnownPartitionSecondaryBootloader,
+	kKnownPartitionSecondaryBootloaderBackup,
+	kKnownPartitionParam,
+	kKnownPartitionKernel,
+	kKnownPartitionRecovery,
+	kKnownPartitionEfs,
+	kKnownPartitionModem,
+
+	kKnownPartitionNormalBoot,
+	kKnownPartitionSystem,
+	kKnownPartitionUserData,
+	kKnownPartitionFota,
+	kKnownPartitionHidden,
+	kKnownPartitionMovinand,
+	kKnownPartitionData,
+	kKnownPartitionUms,
+	kKnownPartitionEmmc,
+
+	kKnownPartitionCount
 };
 
-bool flashFile(BridgeManager *bridgeManager, FILE *file, int fileIndex)
+vector<const char *> knownPartitionNames[kKnownPartitionCount];
+
+struct PartitionNameFilePair
 {
-	switch (fileIndex)
+	string partitionName;
+	FILE *file;
+
+	PartitionNameFilePair(const char *partitionName, FILE *file)
 	{
-		case kFilePit:
+		this->partitionName = partitionName;
+		this->file = file;
+	}
+};
 
-			InterfaceManager::Print("Uploading PIT file\n");
-			if (bridgeManager->SendPitFile(file))
-			{
-				InterfaceManager::Print("PIT file upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("PIT file upload failed!\n");
-				return (false);
-			}
+void initialiseKnownPartitionNames(void)
+{
+	knownPartitionNames[kKnownPartitionPit].push_back("PIT");
+	knownPartitionNames[kKnownPartitionFactoryFs].push_back("FACTORYFS");
+	knownPartitionNames[kKnownPartitionCache].push_back("CACHE");
+	knownPartitionNames[kKnownPartitionDatabaseData].push_back("DBDATAFS");
+
+	knownPartitionNames[kKnownPartitionPrimaryBootloader].push_back("IBL+PBL");
+	knownPartitionNames[kKnownPartitionPrimaryBootloader].push_back("BOOT");
+
+	knownPartitionNames[kKnownPartitionSecondaryBootloader].push_back("SBL");
+	knownPartitionNames[kKnownPartitionSecondaryBootloader].push_back("SBL1");
+
+	knownPartitionNames[kKnownPartitionSecondaryBootloaderBackup].push_back("SBL2");
+	knownPartitionNames[kKnownPartitionParam].push_back("PARAM");
+	knownPartitionNames[kKnownPartitionKernel].push_back("KERNEL");
+	knownPartitionNames[kKnownPartitionRecovery].push_back("RECOVERY");
+	knownPartitionNames[kKnownPartitionEfs].push_back("EFS");
+	knownPartitionNames[kKnownPartitionModem].push_back("MODEM");
+
+	knownPartitionNames[kKnownPartitionNormalBoot].push_back("NORMALBOOT");
+	knownPartitionNames[kKnownPartitionSystem].push_back("SYSTEM");
+	knownPartitionNames[kKnownPartitionUserData].push_back("USERDATA");
+	knownPartitionNames[kKnownPartitionFota].push_back("FOTA");
+	knownPartitionNames[kKnownPartitionHidden].push_back("HIDDEN");
+	knownPartitionNames[kKnownPartitionMovinand].push_back("MOVINAND");
+	knownPartitionNames[kKnownPartitionData].push_back("DATAFS");
+	knownPartitionNames[kKnownPartitionUms].push_back("UMS.EN");
+	knownPartitionNames[kKnownPartitionEmmc].push_back("GANG");
+}
 
-		case kFileFactoryFs:
+bool openFiles(const map<string, string>& argumentMap, map<string, FILE *>& argumentFileMap)
+{
+	map<string, string>::const_iterator it = argumentMap.begin();
 
-			InterfaceManager::Print("Uploading factory filesytem\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileFactoryFilesystem))
-			{
-				InterfaceManager::Print("Factory filesytem upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("Factory filesytem upload failed!\n");
-				return (false);
-			}
+	for (it = argumentMap.begin(); it != argumentMap.end(); it++)
+	{
+		bool isFileArgument = false;
 
-		case kFileCache:
+		int partitionIndex = atoi(it->first.substr(it->first.find_first_not_of('-')).c_str());
 
-			InterfaceManager::Print("Uploading cache\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileCache))
-			{
-				InterfaceManager::Print("Cache upload successful\n");
-				return (true);
-			}
-			else
+		// Was the argument a partition index?
+		if (partitionIndex > 0 || it->first.compare("-0") == 0)
+		{
+			isFileArgument = true;
+		}
+		else
+		{
+			// The argument wasn't a partition index, check if it's a known partition name.
+			for (int knownPartition = 0; knownPartition < kKnownPartitionCount; knownPartition++)
 			{
-				InterfaceManager::PrintError("Cache upload failed!\n");
-				return (false);
+				if (it->first.compare(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit + knownPartition]) == 0)
+				{
+					isFileArgument = true;
+					break;
+				}
 			}
+		}
 
-		case kFileData:
+		if (!isFileArgument)
+			continue;
 
-			InterfaceManager::Print("Uploading data database\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileDatabaseData))
-			{
-				InterfaceManager::Print("Data database upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("Data database upload failed!\n");
-				return (false);
-			}
+		pair<string, FILE *> argumentFilePair;
+		argumentFilePair.first = it->first;
+		argumentFilePair.second = fopen(it->second.c_str(), "rb");
 
-		case kFilePrimaryBootloader:
+		if (!argumentFilePair.second)
+		{
+			InterfaceManager::PrintError("Failed to open file \"%s\"\n", it->second.c_str());
+			return (false);
+		}
 
-			InterfaceManager::Print("Uploading primary bootloader\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFilePrimaryBootloader))
-			{
-				InterfaceManager::Print("Primary bootloader upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("Primary bootloader upload failed!\n");
-				return (false);
-			}
+		argumentFileMap.insert(argumentFilePair);
+	}
 
-		case kFileSecondaryBootloader:
+	return (true);
+}
 
-			InterfaceManager::Print("Uploading secondary bootloader\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileSecondaryBootloader))
-			{
-				InterfaceManager::Print("Secondary bootloader upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("Secondary bootloader upload failed!\n");
-				return (false);
-			}
+bool mapFilesToPartitions(const map<string, FILE *>& argumentFileMap, const PitData *pitData, map<unsigned int, PartitionNameFilePair>& partitionFileMap)
+{
+	map<string, FILE *>::const_iterator it = argumentFileMap.begin();
 
-		case kFileSecondaryBootloaderBackup:
+	for (it = argumentFileMap.begin(); it != argumentFileMap.end(); it++)
+	{
+		int partitionIndex = atoi(it->first.substr(it->first.find_first_not_of('-')).c_str());
 
-			InterfaceManager::Print("Uploading backup secondary bootloader\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileSecondaryBootloaderBackup))
-			{
-				InterfaceManager::Print("Backup secondary bootloader upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("Backup secondary bootloader upload failed!\n");
-				return (false);
-			}
+		const PitEntry *pitEntry = nullptr;
 
-		case kFileParam:
+		// Was the argument a partition index?
+		if (partitionIndex > 0 || it->first.compare("-0") == 0)
+		{
+			pitEntry = pitData->FindEntry(partitionIndex);
+		}
+		else
+		{
+			// The argument wasn't a partition index, so it must be a known partition name.
+			int knownPartition;
 
-			InterfaceManager::Print("Uploading param.lfs\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileParamLfs))
+			for (knownPartition = 0; knownPartition < kKnownPartitionCount; knownPartition++)
 			{
-				InterfaceManager::Print("param.lfs upload successful\n");
-				return (true);
+				if (it->first.compare(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit + knownPartition]) == 0)
+					break;
 			}
-			else
+
+			// Check for the partition in the PIT file using all known names.
+			for (unsigned int i = 0; i < knownPartitionNames[knownPartition].size(); i++)
 			{
-				InterfaceManager::PrintError("param.lfs upload failed!\n");
-				return (false);
+				pitEntry = pitData->FindEntry(knownPartitionNames[knownPartition][i]);
+
+				if (pitEntry)
+					break;
 			}
+		}
 
-		case kFileKernel:
+		if (!pitEntry)
+		{
+			InterfaceManager::PrintError("Partition corresponding to %s argument could not be located\n", it->first.c_str());
+			return (false);
+		}
 
-			InterfaceManager::Print("Uploading kernel\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileKernel))
-			{
-				InterfaceManager::Print("Kernel upload successful\n");
-				return (true);
-			}
-			else
-			{
-				InterfaceManager::PrintError("Kernel upload failed!\n");
-				return (false);
-			}
+		PartitionNameFilePair partitionNameFilePair(pitEntry->GetPartitionName(), it->second);
+		partitionFileMap.insert(pair<unsigned int, PartitionNameFilePair>(pitEntry->GetPartitionIdentifier(), partitionNameFilePair));
+	}
 
-		case kFileModem:
+	return (true);
+}
 
-			InterfaceManager::Print("Uploading modem\n");
-			
-			if (bridgeManager->SendFile(file, EndModemFileTransferPacket::kDestinationModem))     // <-- Odin method
-			/*if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,    // <-- Kies method. WARNING: Doesn't work on Galaxy Tab!
-				EndPhoneFileTransferPacket::kFileModem))*/
-			{
-				InterfaceManager::Print("Modem upload successful\n");
-				return (true);
-			}
-			else
+void closeFiles(map<string, FILE *> argumentfileMap)
+{
+	for (map<string, FILE *>::iterator it = argumentfileMap.begin(); it != argumentfileMap.end(); it++)
+		fclose(it->second);
+
+	argumentfileMap.clear();
+}
+
+int downloadPitFile(BridgeManager *bridgeManager, unsigned char **pitBuffer)
+{
+	InterfaceManager::Print("Downloading device's PIT file...\n");
+
+	int devicePitFileSize = bridgeManager->ReceivePitFile(pitBuffer);
+
+	if (!*pitBuffer)
+	{
+		InterfaceManager::PrintError("Failed to download PIT file!\n");
+
+		return (-1);
+	}
+
+	InterfaceManager::Print("PIT file download sucessful\n\n");
+	return devicePitFileSize;
+}
+
+bool flashFile(BridgeManager *bridgeManager, unsigned int partitionIndex, const char *partitionName, FILE *file)
+{
+	// PIT files need to be handled differently, try determine if the partition we're flashing to is a PIT partition.
+	bool isPit = false;
+
+	for (unsigned int i = 0; i < knownPartitionNames[kKnownPartitionPit].size(); i++)
+	{
+		if (strcmp(partitionName, knownPartitionNames[kKnownPartitionPit][i]) == 0)
+		{
+			isPit = true;
+			break;
+		}
+	}
+
+	if (isPit)
+	{
+		InterfaceManager::Print("Uploading %s\n", partitionName);
+
+		if (bridgeManager->SendPitFile(file))
+		{
+			InterfaceManager::Print("%s upload successful\n", partitionName);
+			return (true);
+		}
+		else
+		{
+			InterfaceManager::PrintError("%s upload failed!\n", partitionName);
+			return (false);
+		}
+	}
+	else
+	{
+		// Modems need to be handled differently, try determine if the partition we're flashing to is a modem partition.
+		bool isModem = false;
+
+		for (unsigned int i = 0; i < knownPartitionNames[kKnownPartitionModem].size(); i++)
+		{
+			if (strcmp(partitionName, knownPartitionNames[kKnownPartitionModem][i]) == 0)
 			{
-				InterfaceManager::PrintError("Modem upload failed!\n");
-				return (false);
+				isModem = true;
+				break;
 			}
+		}
 
-		case kFileRecovery:
+		if (isModem)
+		{			
+			InterfaceManager::Print("Uploading %s\n", partitionName);
 
-			InterfaceManager::Print("Uploading recovery\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileRecovery))
+			//if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,    // <-- Kies method. WARNING: Doesn't work on Galaxy Tab!
+			//	EndPhoneFileTransferPacket::kFileModem))
+			if (bridgeManager->SendFile(file, EndModemFileTransferPacket::kDestinationModem))     // <-- Odin method
 			{
-				InterfaceManager::Print("Recovery upload successful\n");
+				InterfaceManager::Print("%s upload successful\n", partitionName);
 				return (true);
 			}
 			else
 			{
-				InterfaceManager::PrintError("Recovery upload failed!\n");
+				InterfaceManager::PrintError("%s upload failed!\n", partitionName);
 				return (false);
 			}
+		}
+		else
+		{
+			// We're uploading to a phone partition
+			InterfaceManager::Print("Uploading %s\n", partitionName);
 
-		case kFileEfs:
-
-			InterfaceManager::Print("Uploading EFS\n");
-			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone,
-				EndPhoneFileTransferPacket::kFileEfs))
+			if (bridgeManager->SendFile(file, EndPhoneFileTransferPacket::kDestinationPhone, partitionIndex))
 			{
-				InterfaceManager::Print("EFS upload successful\n");
+				InterfaceManager::Print("%s upload successful\n", partitionName);
 				return (true);
 			}
 			else
 			{
-				InterfaceManager::PrintError("EFS upload failed!\n");
+				InterfaceManager::PrintError("%s upload failed!\n", partitionName);
 				return (false);
 			}
-
-		default:
-
-			InterfaceManager::PrintError("ERROR: Attempted to flash unknown file!\n");
-			return (false);
+		}
 	}
+
+	return (true);
 }
 
-bool attemptFlash(BridgeManager *bridgeManager, FILE **fileArray, bool repartition)
+bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFileMap, bool repartition)
 {
 	bool success;
 	
@@ -271,14 +339,10 @@ bool attemptFlash(BridgeManager *bridgeManager, FILE **fileArray, bool repartiti
 		return (false);
 	}
 
-	if (unknown != 0)
+	// 131072 for Galaxy S II, 0 for other devices.
+	if (unknown != 0 && unknown != 131072)
 	{
 		InterfaceManager::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%i\n", unknown);
-
-		if (!bridgeManager->EndSession())
-			return (false);
-		bridgeManager->RebootDevice();
-
 		return (false);
 	}
 
@@ -304,36 +368,20 @@ bool attemptFlash(BridgeManager *bridgeManager, FILE **fileArray, bool repartiti
 		return (false);
 	}
 
-	// TODO: Work out what this value is... it has been either 180 or 0 for Galaxy S phones, and 3 on the Galaxy Tab.
-	if (unknown != 180 && unknown != 0 && unknown != 3)
+	// TODO: Work out what this value is... it has been either 180 or 0 for Galaxy S phones, 3 on the Galaxy Tab, 190 for SHW-M110S.
+	if (unknown != 180 && unknown != 0 && unknown != 3 && unknown != 190)
 	{
 		InterfaceManager::PrintError("Unexpected device info response!\nExpected: 180, 0 or 3\nReceived:%i\n", unknown);
-
-		if (!bridgeManager->EndSession())
-			return (false);
-		bridgeManager->RebootDevice();
-
 		return (false);
 	}
 	// --------------------------------------------------------------
 
 	int totalBytes = 0;
-	for (int i = kFileFactoryFs; i < kFileCount; i++)
-	{
-		if (fileArray[i])
-		{
-			fseek(fileArray[i], 0, SEEK_END);
-			totalBytes += ftell(fileArray[i]);
-			rewind(fileArray[i]);
-		}
-	}
-
-	if (repartition)
+	for (map<string, FILE *>::const_iterator it = argumentFileMap.begin(); it != argumentFileMap.end(); it++)
 	{
-		// When repartitioning we send the PIT file to the device.
-		fseek(fileArray[kFilePit], 0, SEEK_END);
-		totalBytes += ftell(fileArray[kFilePit]);
-		rewind(fileArray[kFilePit]);
+		fseek(it->second, 0, SEEK_END);
+		totalBytes += ftell(it->second);
+		rewind(it->second);
 	}
 	
 	deviceInfoPacket = new DeviceInfoPacket(DeviceInfoPacket::kTotalBytes, totalBytes);
@@ -360,112 +408,91 @@ bool attemptFlash(BridgeManager *bridgeManager, FILE **fileArray, bool repartiti
 	if (unknown != 0)
 	{
 		InterfaceManager::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%i\n", unknown);
-
-		if (!bridgeManager->EndSession())
-			return (false);
-		bridgeManager->RebootDevice();
-
 		return (false);
 	}
 
 	// -----------------------------------------------------
 
-	if (fileArray[kFilePit])
+	PitData *pitData;
+	FILE *localPitFile = nullptr;
+
+	if (repartition)
 	{
-		if (repartition)
+		// If we're repartitioning then we need to unpack the information from the specified PIT file.
+
+		map<string, FILE *>::iterator it = argumentFileMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit]);
+
+		// This shouldn't ever happen due to early checks, but we'll check again just in case...
+		if (it == argumentFileMap.end())
 		{
-			flashFile(bridgeManager, fileArray[kFilePit], kFilePit);
+			InterfaceManager::PrintError("Attempt was made to repartition without specifying a PIT file!\n");
+			return (false);
 		}
-		else // We're performing a PIT check
-		{
-			// Load the local pit file into memory.
-			char *localPit = new char[4096];
-			memset(localPit, 0, 4096);
 
-			fseek(fileArray[kFilePit], 0, SEEK_END);
-			long localPitFileSize = ftell(fileArray[kFilePit]);
-			rewind(fileArray[kFilePit]);
+		localPitFile = it->second;
 
-			fread(localPit, 1, localPitFileSize, fileArray[kFilePit]);
+		// Load the local pit file into memory.
+		unsigned char *pitFileBuffer = new unsigned char[4096];
+		memset(pitFileBuffer, 0, 4096);
 
-			InterfaceManager::Print("Downloading device's PIT file...\n");
+		fseek(localPitFile, 0, SEEK_END);
+		long localPitFileSize = ftell(localPitFile);
+		rewind(localPitFile);
 
-			unsigned char *devicePit;
-			int devicePitFileSize = bridgeManager->ReceivePitFile(&devicePit);
+		fread(pitFileBuffer, 1, localPitFileSize, localPitFile);
+		rewind(localPitFile);
 
-			if (!devicePit)
-			{
-				InterfaceManager::PrintError("Failed to download PIT file!\n");
+		pitData = new PitData();
+		pitData->Unpack(pitFileBuffer);
 
-				if (!bridgeManager->EndSession())
-					return (false);
-				bridgeManager->RebootDevice();
+		delete [] pitFileBuffer;
+	}
+	else
+	{
+		// If we're not repartitioning then we need to retrieve the device's PIT file and unpack it.
 
-				return (false);
-			}
+		unsigned char *pitFileBuffer;
+		downloadPitFile(bridgeManager, &pitFileBuffer);
 
-			InterfaceManager::Print("PIT file download sucessful\n\n");
+		pitData = new PitData();
+		pitData->Unpack(pitFileBuffer);
 
-			bool pitFilesMatch = !memcmp(localPit, devicePit, localPitFileSize);
+		delete [] pitFileBuffer;
+	}
 
-			delete [] localPit;
-			delete [] devicePit;
+	map<unsigned int, PartitionNameFilePair> partitionFileMap;
 
-			if (!pitFilesMatch)
-			{
-				InterfaceManager::Print("Optional PIT check failed! To disable this check don't use the --pit parameter.");
+	// Map the files being flashed to partitions stored in PIT file.
+	mapFilesToPartitions(argumentFileMap, pitData, partitionFileMap);
+	
+	delete pitData;
 
-				if (!bridgeManager->EndSession())
+	// If we're repartitioning then we need to flash the PIT file first.
+	if (repartition)
+	{
+		for (map<unsigned int, PartitionNameFilePair>::iterator it = partitionFileMap.begin(); it != partitionFileMap.end(); it++)
+		{
+			if (it->second.file == localPitFile)
+			{
+				if (!flashFile(bridgeManager, it->first, it->second.partitionName.c_str(), it->second.file))
 					return (false);
-				bridgeManager->RebootDevice();
-
-				return (false);
 			}
 		}
 	}
 
-	// Flash specified files
-	for (int fileIndex = kFileFactoryFs; fileIndex < kFileCount; fileIndex++)
+	// Flash all other files
+	for (map<unsigned int, PartitionNameFilePair>::iterator it = partitionFileMap.begin(); it != partitionFileMap.end(); it++)
 	{
-		if (fileArray[fileIndex])
+		if (it->second.file != localPitFile)
 		{
-			if (!flashFile(bridgeManager, fileArray[fileIndex], fileIndex))
+			if (!flashFile(bridgeManager, it->first, it->second.partitionName.c_str(), it->second.file))
 				return (false);
 		}
 	}
 
-	return (bridgeManager->EndSession() && bridgeManager->RebootDevice());
-}
-
-bool openFiles(const map<string, string>& argumentMap, FILE **fileArray)
-{
-	for (int fileIndex = 0; fileIndex < kFileCount; fileIndex++)
-	{
-		// kFlashArgPit + kFile<Name> == kFlashArg<Name>
-		map<string, string>::const_iterator it = argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit + fileIndex]);
-		if (it == argumentMap.end())
-			continue;
-
-		fileArray[fileIndex] = fopen(it->second.c_str(), "rb");
-		if (!fileArray[fileIndex])
-		{
-			InterfaceManager::PrintError("Failed to open file \"%s\"\n", it->second.c_str());
-			return (false);
-		}
-	}
-
 	return (true);
 }
 
-void closeFiles(FILE **fileArray)
-{
-	for (int fileIndex = 0; fileIndex < kFileCount; fileIndex++)
-	{
-		if (fileArray[fileIndex] != nullptr)
-			fclose(fileArray[fileIndex]);
-	}
-}
-
 int main(int argc, char **argv)
 {
 	map<string, string> argumentMap;
@@ -473,10 +500,12 @@ int main(int argc, char **argv)
 
 	if (!InterfaceManager::GetArguments(argc, argv, argumentMap, &actionIndex))
 	{
-		Sleep(1000);
+		Sleep(250);
 		return (0);
 	}
 
+	initialiseKnownPartitionNames();
+
 	if (actionIndex == InterfaceManager::kActionHelp)
 	{
 		InterfaceManager::Print(InterfaceManager::usage);
@@ -485,16 +514,16 @@ int main(int argc, char **argv)
 	else if (actionIndex == InterfaceManager::kActionFlash)
 	{
 		if (argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgRepartition]) != argumentMap.end()
-			&& (argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgFactoryFs]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgCache]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgData]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPrimaryBootloader]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgSecondaryBootloader]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgParam]) == argumentMap.end()
-			|| argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgKernel]) == argumentMap.end()))
+			&& argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit]) == argumentMap.end())
 		{
-			InterfaceManager::Print("If you wish to repartition then factoryfs, cache, dbdata, primary and secondary\nbootloaders, param, kernel and a PIT file must all be specified.\n");
+			InterfaceManager::Print("If you wish to repartition then a PIT file must be specified.\n");
+			return (0);
+		}
+
+		if (argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgPit]) != argumentMap.end()
+			&& argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgRepartition]) == argumentMap.end())
+		{
+			InterfaceManager::Print("A PIT file should only be used when repartitioning.\n");
 			return (0);
 		}
 	}
@@ -537,7 +566,7 @@ int main(int argc, char **argv)
 		}
 	}
 
-	InterfaceManager::Print("\nHeimdall, Copyright (c) 2010, Benjamin Dobell, Glass Echidna\n");
+	InterfaceManager::Print("\nHeimdall v, Copyright (c) 2010, Benjamin Dobell, Glass Echidna\n");
 	InterfaceManager::Print("http://www.glassechidna.com.au\n\n");
 	InterfaceManager::Print("This software is provided free of charge. Copying and redistribution is\nencouraged.\n\n");
 	InterfaceManager::Print("If you appreciate this software and you would like to support future\ndevelopment please consider donating:\n");
@@ -546,6 +575,7 @@ int main(int argc, char **argv)
 	Sleep(1000);
 
 	bool verbose = argumentMap.find(InterfaceManager::commonArgumentNames[InterfaceManager::kCommonArgVerbose]) != argumentMap.end();
+	bool noReboot = argumentMap.find(InterfaceManager::commonArgumentNames[InterfaceManager::kCommonArgNoReboot]) != argumentMap.end();
 
 	int communicationDelay = BridgeManager::kCommunicationDelayDefault;
 	if (argumentMap.find(InterfaceManager::commonArgumentNames[InterfaceManager::kCommonArgDelay]) != argumentMap.end())
@@ -565,16 +595,12 @@ int main(int argc, char **argv)
 	{
 		case InterfaceManager::kActionFlash:
 		{
-			FILE **fileArray = new FILE *[kFileCount];
-			for (int i = 0; i < kFileCount; i++)
-				fileArray[i] = nullptr;
+			map<string, FILE *> argumentFileMap;
 
 			// We open the files before doing anything else to ensure they exist.
-			if (!openFiles(argumentMap, fileArray))
+			if (!openFiles(argumentMap, argumentFileMap))
 			{
-				closeFiles(fileArray);
-				delete [] fileArray;
-
+				closeFiles(argumentFileMap);
 				delete bridgeManager;
 
 				return (0);
@@ -582,19 +608,21 @@ int main(int argc, char **argv)
 
 			if (!bridgeManager->BeginSession())
 			{
-				closeFiles(fileArray);
-				delete [] fileArray;
-
+				closeFiles(argumentFileMap);
 				delete bridgeManager;
 
 				return (-1);
 			}
 
 			bool repartition = argumentMap.find(InterfaceManager::flashArgumentNames[InterfaceManager::kFlashArgRepartition]) != argumentMap.end();
-			success = attemptFlash(bridgeManager, fileArray, repartition);
+			success = attemptFlash(bridgeManager, argumentFileMap, repartition);
+
+			if (noReboot)
+				success = bridgeManager->EndSession() && success;
+			else
+				success = bridgeManager->EndSession() && bridgeManager->RebootDevice() && success;
 
-			closeFiles(fileArray);
-			delete [] fileArray;
+			closeFiles(argumentFileMap);
 
 			break;
 		}
@@ -608,7 +636,11 @@ int main(int argc, char **argv)
 			}
 
 			InterfaceManager::Print("Attempting to close connect to pc screen...\n");
-			success = bridgeManager->EndSession() && bridgeManager->RebootDevice();
+
+			if (noReboot)
+				success = bridgeManager->EndSession();
+			else
+				success = bridgeManager->EndSession() && bridgeManager->RebootDevice();
 
 			if (success)
 				InterfaceManager::Print("Attempt complete\n");
@@ -640,7 +672,6 @@ int main(int argc, char **argv)
 				fclose(dumpFile);
 
 				delete bridgeManager;
-
 				return (-1);
 			}
 
@@ -648,8 +679,55 @@ int main(int argc, char **argv)
 
 			fclose(dumpFile);
 
-			if (success)
-				success = bridgeManager->EndSession() && bridgeManager->RebootDevice();
+			if (noReboot)
+				success = bridgeManager->EndSession() && success;
+			else
+				success = bridgeManager->EndSession() && bridgeManager->RebootDevice() && success;
+
+			break;
+		}
+
+		case InterfaceManager::kActionPrintPit:
+		{
+			if (!bridgeManager->BeginSession())
+			{
+				delete bridgeManager;
+				return (-1);
+			}
+
+			unsigned char *devicePit;
+
+			if (downloadPitFile(bridgeManager, &devicePit) < -1)
+			{
+				if (!bridgeManager->EndSession())
+					return (-1);
+
+				if (!noReboot)
+					bridgeManager->RebootDevice();
+
+				delete bridgeManager;
+				return (-1);
+			}
+
+			PitData *pitData = new PitData();
+
+			if (pitData->Unpack(devicePit))
+			{
+				pitData->Print();
+				success = true;
+			}
+			else
+			{
+				InterfaceManager::PrintError("Failed to unpack device's PIT file!\n");
+				success = false;
+			}
+
+			delete pitData;
+
+			if (noReboot)
+				success = bridgeManager->EndSession() && success;
+			else
+				success = bridgeManager->EndSession() && bridgeManager->RebootDevice() && success;
 
 			break;
 		}