4 files changed, 61 insertions, 44 deletions

diff --git a/heimdall/source/BridgeManager.cpp b/heimdall/source/BridgeManager.cpp
index 204d87f..ef5faa6 100644
--- a/heimdall/source/BridgeManager.cpp
+++ b/heimdall/source/BridgeManager.cpp
@@ -116,7 +116,7 @@ bool BridgeManager::DetectDevice(void)
 	int result = libusb_init(&libusbContext);
 	if (result != LIBUSB_SUCCESS)
 	{
-		Interface::PrintError("Failed to initialise libusb. Error: %i\n", result);
+		Interface::PrintError("Failed to initialise libusb. Error: %d\n", result);
 		return (false);
 	}
 
@@ -153,7 +153,7 @@ bool BridgeManager::Initialise(void)
 	int result = libusb_init(&libusbContext);
 	if (result != LIBUSB_SUCCESS)
 	{
-		Interface::PrintError("Failed to initialise libusb. Error: %i\n", result);
+		Interface::PrintError("Failed to initialise libusb. Error: %d\n", result);
 		return (false);
 	}
 
@@ -191,7 +191,7 @@ bool BridgeManager::Initialise(void)
 	result = libusb_open(heimdallDevice, &deviceHandle);
 	if (result != LIBUSB_SUCCESS)
 	{
-		Interface::PrintError("Failed to access device. Error: %i\n", result);
+		Interface::PrintError("Failed to access device. Error: %d\n", result);
 		return (false);
 	}
 
@@ -528,7 +528,7 @@ bool BridgeManager::SendPacket(OutboundPacket *packet, int timeout) const
 		int retryDelay = (communicationDelay > 250) ? communicationDelay : 250;
 
 		if (verbose)
-			Interface::PrintError("Error %i whilst sending packet. ", result);
+			Interface::PrintError("Error %d whilst sending packet. ", result);
 
 		// Retry
 		for (int i = 0; i < 5; i++)
@@ -546,7 +546,7 @@ bool BridgeManager::SendPacket(OutboundPacket *packet, int timeout) const
 				break;
 
 			if (verbose)
-				Interface::PrintError("Error %i whilst sending packet. ", result);
+				Interface::PrintError("Error %d whilst sending packet. ", result);
 		}
 
 		if (verbose)
@@ -574,7 +574,7 @@ bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout) const
 		int retryDelay = (communicationDelay > 250) ? communicationDelay : 250;
 
 		if (verbose)
-			Interface::PrintError("Error %i whilst receiving packet. ", result);
+			Interface::PrintError("Error %d whilst receiving packet. ", result);
 
 		// Retry
 		for (int i = 0; i < 5; i++)
@@ -592,7 +592,7 @@ bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout) const
 				break;
 
 			if (verbose)
-				Interface::PrintError("Error %i whilst receiving packet. ", result);
+				Interface::PrintError("Error %d whilst receiving packet. ", result);
 
 			if (i >= 3)
 			{
@@ -750,7 +750,7 @@ int BridgeManager::ReceivePitFile(unsigned char **pitBuffer) const
 
 		if (!success)
 		{
-			Interface::PrintError("Failed to request PIT file part #%i!\n", i);
+			Interface::PrintError("Failed to request PIT file part #%d!\n", i);
 			delete [] buffer;
 			return (0);
 		}
@@ -760,7 +760,7 @@ int BridgeManager::ReceivePitFile(unsigned char **pitBuffer) const
 		
 		if (!success)
 		{
-			Interface::PrintError("Failed to receive PIT file part #%i!\n", i);
+			Interface::PrintError("Failed to receive PIT file part #%d!\n", i);
 			delete receiveFilePartPacket;
 			delete [] buffer;
 			return (0);
@@ -906,7 +906,7 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 			if (verbose)
 			{
 				const unsigned char *data = sendFilePartResponse->GetData();
-				Interface::Print("File Part #%i... Response: %X  %X  %X  %X  %X  %X  %X  %X \n", filePartIndex,
+				Interface::Print("File Part #%d... Response: %X  %X  %X  %X  %X  %X  %X  %X \n", filePartIndex,
 					data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7]);
 			}
 
@@ -939,7 +939,7 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 					if (verbose)
 					{
 						const unsigned char *data = sendFilePartResponse->GetData();
-						Interface::Print("File Part #%i... Response: %X  %X  %X  %X  %X  %X  %X  %X \n", filePartIndex,
+						Interface::Print("File Part #%d... Response: %X  %X  %X  %X  %X  %X  %X  %X \n", filePartIndex,
 							data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7]);
 					}
 
@@ -947,7 +947,7 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 					if (receivedPartIndex != filePartIndex)
 					{
-						Interface::PrintError("\nERROR: Expected file part index: %i Received: %i\n",
+						Interface::PrintError("\nERROR: Expected file part index: %d Received: %d\n",
 							filePartIndex, receivedPartIndex);
 						return (false);
 					}
@@ -962,7 +962,7 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 			if (receivedPartIndex != filePartIndex)
 			{
-				Interface::PrintError("\nERROR: Expected file part index: %i Received: %i\n",
+				Interface::PrintError("\nERROR: Expected file part index: %d Received: %d\n",
 					filePartIndex, receivedPartIndex);
 				return (false);
 			}
@@ -978,9 +978,9 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 				if (currentPercent != previousPercent)
 				{
 					if (previousPercent < 10)
-						Interface::Print("\b\b%i%%", currentPercent);
+						Interface::Print("\b\b%d%%", currentPercent);
 					else
-						Interface::Print("\b\b\b%i%%", currentPercent);
+						Interface::Print("\b\b\b%d%%", currentPercent);
 				}
 			}
 
@@ -1076,7 +1076,7 @@ bool BridgeManager::ReceiveDump(int chipType, int chipId, FILE *file) const
 
 		if (!success)
 		{
-			Interface::PrintError("Failed to request dump part #%i!\n", i);
+			Interface::PrintError("Failed to request dump part #%d!\n", i);
 			delete [] buffer;
 			return (false);
 		}
@@ -1086,7 +1086,7 @@ bool BridgeManager::ReceiveDump(int chipType, int chipId, FILE *file) const
 		
 		if (!success)
 		{
-			Interface::PrintError("Failed to receive dump part #%i!\n", i);
+			Interface::PrintError("Failed to receive dump part #%d!\n", i);
 			continue;
 			delete receiveFilePartPacket;
 			delete [] buffer;
diff --git a/heimdall/source/Interface.cpp b/heimdall/source/Interface.cpp
index 5f0942d..35105bb 100644
--- a/heimdall/source/Interface.cpp
+++ b/heimdall/source/Interface.cpp
@@ -31,7 +31,7 @@ using namespace std;
 using namespace libpit;
 using namespace Heimdall;
 
-const char *Interface::version = "v1.3.0";
+const char *Interface::version = "v1.3 (beta)";
 
 const char *Interface::usage = "Usage: heimdall <action> <arguments> [--verbose] [--no-reboot] [--delay <ms>]\n\
 \n\
diff --git a/heimdall/source/Interface.h b/heimdall/source/Interface.h
index c6bdef3..7b1741d 100644
--- a/heimdall/source/Interface.h
+++ b/heimdall/source/Interface.h
@@ -249,6 +249,11 @@ namespace Heimdall
 			static void PrintReleaseInfo(void);
 
 			static void PrintPit(const PitData *pitData);
+
+			static string& GetPitArgument(void)
+			{
+				return (flashValueArguments[kFlashValueArgPit]);
+			}
 	};
 }
 
diff --git a/heimdall/source/main.cpp b/heimdall/source/main.cpp
index 64cc553..f24b64a 100644
--- a/heimdall/source/main.cpp
+++ b/heimdall/source/main.cpp
@@ -348,7 +348,7 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 	// 131072 for Galaxy S II, 0 for other devices.
 	if (deviceInfoResult != 0 && deviceInfoResult != 131072)
 	{
-		Interface::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%i\n", deviceInfoResult);
+		Interface::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%d\n", deviceInfoResult);
 		return (false);
 	}
 
@@ -360,7 +360,7 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 	// 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 (deviceInfoResult != 180 && deviceInfoResult != 0 && deviceInfoResult != 3 && deviceInfoResult != 190)
 	{
-		Interface::PrintError("Unexpected device info response!\nExpected: 180, 0 or 3\nReceived:%i\n", deviceInfoResult);
+		Interface::PrintError("Unexpected device info response!\nExpected: 180, 0 or 3\nReceived:%d\n", deviceInfoResult);
 		return (false);
 	}
 
@@ -369,9 +369,12 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 	int totalBytes = 0;
 	for (map<string, FILE *>::const_iterator it = argumentFileMap.begin(); it != argumentFileMap.end(); it++)
 	{
-		fseek(it->second, 0, SEEK_END);
-		totalBytes += ftell(it->second);
-		rewind(it->second);
+		if (repartition || it->first != Interface::GetPitArgument())
+		{
+			fseek(it->second, 0, SEEK_END);
+			totalBytes += ftell(it->second);
+			rewind(it->second);
+		}
 	}
 	
 	DeviceInfoPacket *deviceInfoPacket = new DeviceInfoPacket(DeviceInfoPacket::kTotalBytes, totalBytes);
@@ -397,27 +400,22 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 
 	if (deviceInfoResult != 0)
 	{
-		Interface::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%i\n", deviceInfoResult);
+		Interface::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%d\n", deviceInfoResult);
 		return (false);
 	}
 
 	// -----------------------------------------------------
 
 	PitData *pitData;
-	FILE *localPitFile = nullptr;
+	PitData *localPitData = nullptr;
 
-	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(Interface::actions[Interface::kActionFlash].valueArguments[Interface::kFlashValueArgPit]);
+	FILE *localPitFile = nullptr;
 
-		// This shouldn't ever happen due to early checks, but we'll check again just in case...
-		if (it == argumentFileMap.end())
-		{
-			Interface::PrintError("Attempt was made to repartition without specifying a PIT file!\n");
-			return (false);
-		}
+	// If a PIT file was passed as an argument then we must unpack it.
+	map<string, FILE *>::iterator it = argumentFileMap.find(Interface::actions[Interface::kActionFlash].valueArguments[Interface::kFlashValueArgPit]);
 
+	if (it != argumentFileMap.end())
+	{
 		localPitFile = it->second;
 
 		// Load the local pit file into memory.
@@ -432,11 +430,17 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 		int dataRead = fread(pitFileBuffer, 1, localPitFileSize, localPitFile);
 		rewind(localPitFile);
 
-		pitData = new PitData();
-		pitData->Unpack(pitFileBuffer);
+		localPitData = new PitData();
+		localPitData->Unpack(pitFileBuffer);
 
 		delete [] pitFileBuffer;
 	}
+	
+	if (repartition)
+	{
+		// Use the local PIT file data.
+		pitData = localPitData;
+	}
 	else
 	{
 		// If we're not repartitioning then we need to retrieve the device's PIT file and unpack it.
@@ -447,6 +451,21 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 		pitData->Unpack(pitFileBuffer);
 
 		delete [] pitFileBuffer;
+
+		if (localPitData != nullptr)
+		{
+			// The user has specified a PIT without repartitioning, we should verify the local and device PIT data match!
+			bool pitsMatch = pitData->Matches(localPitData);
+			delete localPitData;
+
+			if (!pitsMatch)
+			{
+				Interface::PrintError("Local and device PIT files don't match and repartition wasn't specified!\n");
+				
+				delete pitData;
+				return (false);
+			}
+		}
 	}
 
 	map<unsigned int, PartitionNameFilePair> partitionFileMap;
@@ -521,13 +540,6 @@ int main(int argc, char **argv)
 				return (0);
 			}
 
-			if (argumentMap.find(Interface::actions[Interface::kActionFlash].valueArguments[Interface::kFlashValueArgPit]) != argumentMap.end()
-				&& argumentMap.find(Interface::actions[Interface::kActionFlash].valuelessArguments[Interface::kFlashValuelessArgRepartition]) == argumentMap.end())
-			{
-				Interface::Print("A PIT file should only be used when repartitioning.\n");
-				return (0);
-			}
-
 			break;
 
 		case Interface::kActionDump: