4 files changed, 171 insertions, 24 deletions

diff --git a/heimdall/source/BridgeManager.cpp b/heimdall/source/BridgeManager.cpp
index af968a9..f58c4e2 100644
--- a/heimdall/source/BridgeManager.cpp
+++ b/heimdall/source/BridgeManager.cpp
@@ -885,7 +885,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 		if (!success)
 		{
-			Interface::PrintError("\nERROR: Failed to begin file transfer sequence!\n");
+			Interface::PrintErrorSameLine("\n");
+			Interface::PrintError("Failed to begin file transfer sequence!\n");
 			return (false);
 		}
 
@@ -895,7 +896,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 		if (!success)
 		{
-			Interface::PrintError("\nERROR: Failed to confirm beginning of file transfer sequence!\n");
+			Interface::PrintErrorSameLine("\n");
+			Interface::PrintError("Failed to confirm beginning of file transfer sequence!\n");
 			return (false);
 		}
 
@@ -911,7 +913,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 			if (!success)
 			{
-				Interface::PrintError("\nERROR: Failed to send file part packet!\n");
+				Interface::PrintErrorSameLine("\n");
+				Interface::PrintError("Failed to send file part packet!\n");
 				return (false);
 			}
 
@@ -931,11 +934,13 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 			if (!success)
 			{
-				Interface::PrintError("\nERROR: Failed to receive file part response!\n");
+				Interface::PrintErrorSameLine("\n");
+				Interface::PrintError("Failed to receive file part response!\n");
 
 				for (int retry = 0; retry < 4; retry++)
 				{
-					Interface::PrintError("\nERROR: Retrying...");
+					Interface::PrintErrorSameLine("\n");
+					Interface::PrintError("Retrying...");
 
 					// Send
 					sendFilePartPacket = new SendFilePartPacket(file);
@@ -944,7 +949,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 					if (!success)
 					{
-						Interface::PrintError("\nERROR: Failed to send file part packet!\n");
+						Interface::PrintErrorSameLine("\n");
+						Interface::PrintError("Failed to send file part packet!\n");
 						return (false);
 					}
 
@@ -964,8 +970,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 					if (receivedPartIndex != filePartIndex)
 					{
-						Interface::PrintError("\nERROR: Expected file part index: %d Received: %d\n",
-							filePartIndex, receivedPartIndex);
+						Interface::PrintErrorSameLine("\n");
+						Interface::PrintError("Expected file part index: %d Received: %d\n", filePartIndex, receivedPartIndex);
 						return (false);
 					}
 
@@ -979,8 +985,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 			if (receivedPartIndex != filePartIndex)
 			{
-				Interface::PrintError("\nERROR: Expected file part index: %d Received: %d\n",
-					filePartIndex, receivedPartIndex);
+				Interface::PrintErrorSameLine("\n");
+				Interface::PrintError("Expected file part index: %d Received: %d\n", filePartIndex, receivedPartIndex);
 				return (false);
 			}
 
@@ -1021,7 +1027,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 			if (!success)
 			{
-				Interface::PrintError("\nERROR: Failed to end phone file transfer sequence!\n");
+				Interface::PrintErrorSameLine("\n");
+				Interface::PrintError("Failed to end phone file transfer sequence!\n");
 				return (false);
 			}
 		}
@@ -1035,7 +1042,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 			if (!success)
 			{
-				Interface::PrintError("\nERROR: Failed to end modem file transfer sequence!\n");
+				Interface::PrintErrorSameLine("\n");
+				Interface::PrintError("Failed to end modem file transfer sequence!\n");
 				return (false);
 			}
 		}
@@ -1046,7 +1054,8 @@ bool BridgeManager::SendFile(FILE *file, int destination, int fileIdentifier) co
 
 		if (!success)
 		{
-			Interface::PrintError("\nERROR: Failed to confirm end of file transfer sequence!\n");
+			Interface::PrintErrorSameLine("\n");
+			Interface::PrintError("Failed to confirm end of file transfer sequence!\n");
 			return (false);
 		}
 	}
diff --git a/heimdall/source/Interface.cpp b/heimdall/source/Interface.cpp
index 6c3c58e..bb0f30b 100644
--- a/heimdall/source/Interface.cpp
+++ b/heimdall/source/Interface.cpp
@@ -36,9 +36,11 @@ bool Interface::stdoutErrors = false;
 const char *Interface::version = "v1.3 (beta)";
 
 const char *Interface::usage = "Usage: heimdall <action> <action arguments> <common arguments>\n\
+\n\
 Common Arguments:\n\
     [--verbose] [--no-reboot] [--stdout-errors] [--delay <ms>]\n\
 \n\
+\n\
 Action: flash\n\
 Arguments:\n\
     --repartition --pit <filename> [--factoryfs <filename>]\n\
@@ -64,6 +66,11 @@ WARNING: If you're repartitioning it's strongly recommended you specify\n\
 Action: close-pc-screen\n\
 Description: Attempts to get rid off the \"connect phone to PC\" screen.\n\
 \n\
+Action: download-pit\n\
+Arguments: --output <filename>\n\
+Description: Downloads the connected device's PIT file to the specified\n\
+    output file.\n\
+\n\
 Action: detect\n\
 Description: Indicates whether or not a download mode device can be detected.\n\
 \n\
@@ -77,7 +84,7 @@ 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 version\n\
+Action: version\n\
 Description: Displays the version number of this binary.\n\
 \n\
 Action: help\n\
@@ -89,6 +96,12 @@ This software is provided free of charge. Copying and redistribution is\nencoura
 If you appreciate this software and you would like to support future\ndevelopment please consider donating:\n\
 http://www.glassechidna.com.au/donate/\n\n";
 
+const char *Interface::extraInfo = "Heimdall utilises libusb-1.0 for all USB communication:\n\
+    http://www.libusb.org/\n\
+\n\
+libusb-1.0 is licensed under the LGPL-2.1:\n\
+    http://www.gnu.org/licenses/licenses.html#LGPL\n\n";
+
 // Flash arguments
 string Interface::flashValueArguments[kFlashValueArgCount] = {
 	"-pit", "-factoryfs", "-cache", "-dbdata", "-primary-boot",	"-secondary-boot", "-secondary-boot-backup", "-param", "-kernel", "-recovery", "-efs", "-modem",
@@ -108,6 +121,15 @@ string Interface::flashValuelessShortArguments[kFlashValuelessArgCount] = {
 	"r"
 };
 
+// Download PIT arguments
+string Interface::downloadPitValueArguments[kDownloadPitValueArgCount] = {
+	"-output"
+};
+
+string Interface::downloadPitValueShortArguments[kDownloadPitValueArgCount] = {
+	"o"
+};
+
 // Dump arguments
 string Interface::dumpValueArguments[kDumpValueArgCount] = {
 	"-chip-type", "-chip-id", "-output"
@@ -161,7 +183,15 @@ Action Interface::actions[Interface::kActionCount] = {
 
 	// kActionDetect
 	Action("detect", nullptr, nullptr, kDetectValueArgCount,
-		nullptr, nullptr, kDetectValuelessArgCount)
+		nullptr, nullptr, kDetectValuelessArgCount),
+
+	// kActionDownloadPit
+	Action("download-pit", downloadPitValueArguments, downloadPitValueShortArguments, kDownloadPitValueArgCount,
+		nullptr, nullptr, kDownloadPitValuelessArgCount),
+
+	// kActionInfo
+	Action("info", nullptr, nullptr, kInfoValueArgCount,
+		nullptr, nullptr, kInfoValuelessArgCount)
 };
 
 bool Interface::GetArguments(int argc, char **argv, map<string, string>& argumentMap, int *actionIndex)
@@ -383,6 +413,12 @@ void Interface::PrintReleaseInfo(void)
 	Print(releaseInfo, version);
 }
 
+void Interface::PrintFullInfo(void)
+{
+	Print(releaseInfo, version);
+	Print(extraInfo);
+}
+
 void Interface::PrintPit(const PitData *pitData)
 {
 	Interface::Print("Entry Count: %d\n", pitData->GetEntryCount());
diff --git a/heimdall/source/Interface.h b/heimdall/source/Interface.h
index 81f17e5..338445b 100644
--- a/heimdall/source/Interface.h
+++ b/heimdall/source/Interface.h
@@ -76,6 +76,8 @@ namespace Heimdall
 				kActionVersion,
 				kActionHelp,
 				kActionDetect,
+				kActionDownloadPit,
+				kActionInfo,
 				kActionCount
 			};
 
@@ -182,6 +184,18 @@ namespace Heimdall
 				kHelpValuelessArgCount = 0
 			};
 
+			// Info value arguments
+			enum
+			{
+				kInfoValueArgCount = 0
+			};
+
+			// Info valueless arguments
+			enum
+			{
+				kInfoValuelessArgCount = 0
+			};
+
 			// Detect value arguments
 			enum
 			{
@@ -194,6 +208,19 @@ namespace Heimdall
 				kDetectValuelessArgCount = 0
 			};
 
+			// Download PIT value arguments
+			enum
+			{
+				kDownloadPitValueArgOutput = 0,
+				kDownloadPitValueArgCount
+			};
+
+			// Download PIT valueless arguments
+			enum
+			{
+				kDownloadPitValuelessArgCount = 0
+			};
+
 			// Common value arguments
 			enum
 			{
@@ -218,7 +245,8 @@ namespace Heimdall
 		
 			static const char *version;
 			static const char *usage;
-			static const char *releaseInfo;		
+			static const char *releaseInfo;
+			static const char *extraInfo;
 
 			// Flash arguments
 			static string flashValueArguments[kFlashValueArgCount];
@@ -227,6 +255,10 @@ namespace Heimdall
 			static string flashValuelessArguments[kFlashValuelessArgCount];
 			static string flashValuelessShortArguments[kFlashValuelessArgCount];
 
+			// Download PIT arguments
+			static string downloadPitValueArguments[kDownloadPitValueArgCount];
+			static string downloadPitValueShortArguments[kDownloadPitValueArgCount];
+
 			// Dump arguments
 			static string dumpValueArguments[kDumpValueArgCount];
 			static string dumpValueShortArguments[kDumpValueArgCount];
@@ -251,6 +283,7 @@ namespace Heimdall
 			static void PrintVersion(void);
 			static void PrintUsage(void);
 			static void PrintReleaseInfo(void);
+			static void PrintFullInfo(void);
 
 			static void PrintPit(const PitData *pitData);
 
diff --git a/heimdall/source/main.cpp b/heimdall/source/main.cpp
index 7e46c19..63056bd 100644
--- a/heimdall/source/main.cpp
+++ b/heimdall/source/main.cpp
@@ -213,6 +213,14 @@ bool mapFilesToPartitions(const map<string, FILE *>& argumentFileMap, const PitD
 				if (pitEntry)
 					break;
 			}
+
+			if (!pitEntry && knownPartition == kKnownPartitionPit)
+			{
+				PartitionNameFilePair partitionNameFilePair(knownPartitionNames[kKnownPartitionPit][0], it->second);
+				partitionFileMap.insert(pair<unsigned int, PartitionNameFilePair>(static_cast<unsigned int>(-1), partitionNameFilePair));
+
+				return (true);
+			}
 		}
 
 		if (!pitEntry)
@@ -250,7 +258,7 @@ int downloadPitFile(BridgeManager *bridgeManager, unsigned char **pitBuffer)
 	}
 
 	Interface::Print("PIT file download sucessful\n\n");
-	return devicePitFileSize;
+	return (devicePitFileSize);
 }
 
 bool flashFile(BridgeManager *bridgeManager, unsigned int partitionIndex, const char *partitionName, FILE *file)
@@ -537,7 +545,18 @@ int main(int argc, char **argv)
 			if (argumentMap.find(Interface::actions[Interface::kActionFlash].valuelessArguments[Interface::kFlashValuelessArgRepartition]) != argumentMap.end()
 				&& argumentMap.find(Interface::actions[Interface::kActionFlash].valueArguments[Interface::kFlashValueArgPit]) == argumentMap.end())
 			{
-				Interface::Print("If you wish to repartition then a PIT file must be specified.\n");
+				Interface::Print("If you wish to repartition then a PIT file must be specified.\n\n");
+				Interface::PrintUsage();
+				return (0);
+			}
+
+			break;
+
+		case Interface::kActionDownloadPit:
+			if (argumentMap.find(Interface::actions[Interface::kActionDownloadPit].valueArguments[Interface::kDownloadPitValueArgOutput]) == argumentMap.end())
+			{
+				Interface::Print("Output file was not specified.\n\n");
+				Interface::PrintUsage();
 				return (0);
 			}
 
@@ -547,7 +566,7 @@ int main(int argc, char **argv)
 		{
 			if (argumentMap.find(Interface::actions[Interface::kActionDump].valueArguments[Interface::kDumpValueArgOutput]) == argumentMap.end())
 			{
-				Interface::Print("Output file not specified.\n\n");
+				Interface::Print("Output file was not specified.\n\n");
 				Interface::PrintUsage();
 				return (0);
 			}
@@ -578,6 +597,7 @@ int main(int argc, char **argv)
 			if (chipId < 0)
 			{
 				Interface::Print("Chip ID must be a non-negative integer.\n");
+				Interface::PrintUsage();
 				return (0);
 			}
 
@@ -591,6 +611,10 @@ int main(int argc, char **argv)
 		case Interface::kActionHelp:
 			Interface::PrintUsage();
 			return (0);
+
+		case Interface::kActionInfo:
+			Interface::PrintFullInfo();
+			return (0);
 	}
 
 	bool verbose = argumentMap.find(Interface::commonValuelessArguments[Interface::kCommonValuelessArgVerbose]) != argumentMap.end();
@@ -599,6 +623,7 @@ int main(int argc, char **argv)
 	Interface::SetStdoutErrors(argumentMap.find(Interface::commonValuelessArguments[Interface::kCommonValuelessArgStdoutErrors]) != argumentMap.end());
 
 	int communicationDelay = BridgeManager::kCommunicationDelayDefault;
+
 	if (argumentMap.find(Interface::commonValueArguments[Interface::kCommonValueArgDelay]) != argumentMap.end())
 		communicationDelay = atoi(argumentMap.find(Interface::commonValueArguments[Interface::kCommonValueArgDelay])->second.c_str());
 
@@ -606,10 +631,10 @@ int main(int argc, char **argv)
 
 	if (actionIndex == Interface::kActionDetect)
 	{
-		bridgeManager->DetectDevice();
-
+		bool detected = bridgeManager->DetectDevice();
 		delete bridgeManager;
-		return (0);
+
+		return ((detected) ? 0 : 1);
 	}
 
 	Interface::PrintReleaseInfo();
@@ -618,7 +643,7 @@ int main(int argc, char **argv)
 	if (!bridgeManager->Initialise())
 	{
 		delete bridgeManager;
-		return (-2);
+		return (0);
 	}
 
 	bool success;
@@ -674,6 +699,49 @@ int main(int argc, char **argv)
 			break;
 		}
 
+		case Interface::kActionDownloadPit:
+		{
+			map<string, string>::const_iterator it = argumentMap.find(Interface::actions[Interface::kActionDownloadPit].valueArguments[Interface::kDownloadPitValueArgOutput]);
+			FILE *outputPitFile = fopen(it->second.c_str(), "wb");
+
+			if (!outputPitFile)
+			{
+				delete bridgeManager;
+				return (0);
+			}
+
+			if (!bridgeManager->BeginSession())
+			{
+				delete bridgeManager;
+				fclose(outputPitFile);
+				return (-1);
+			}
+
+			unsigned char *pitBuffer;
+			int fileSize = downloadPitFile(bridgeManager, &pitBuffer);
+
+			if (fileSize > 0)
+			{
+				success = fwrite(pitBuffer, 1, fileSize, outputPitFile) == fileSize;
+				fclose(outputPitFile);
+
+				if (!success)
+					Interface::PrintError("Failed to write PIT data to output file.\n");
+
+				success = bridgeManager->EndSession(reboot) && success;
+			}
+			else
+			{
+				fclose(outputPitFile);
+				success = false;
+				bridgeManager->EndSession(reboot);
+			}
+
+			delete [] pitBuffer;
+
+			break;
+		}
+
 		case Interface::kActionDump:
 		{
 			const char *outputFilename = argumentMap.find(Interface::actions[Interface::kActionDump].valueArguments[Interface::kDumpValueArgOutput])->second.c_str();
@@ -740,7 +808,8 @@ int main(int argc, char **argv)
 				Interface::PrintError("Failed to unpack device's PIT file!\n");
 				success = false;
 			}
-
+			
+			delete [] devicePit;
 			delete pitData;
 
 			success = bridgeManager->EndSession(reboot) && success;