1 files changed, 43 insertions, 2 deletions

diff --git a/heimdall/source/DownloadPitAction.cpp b/heimdall/source/DownloadPitAction.cpp
index b4b81a9..744ba56 100644
--- a/heimdall/source/DownloadPitAction.cpp
+++ b/heimdall/source/DownloadPitAction.cpp
@@ -28,11 +28,12 @@
 #include "Heimdall.h"
 #include "Interface.h"
 
+using namespace std;
 using namespace Heimdall;
 
 const char *DownloadPitAction::usage = "Action: download-pit\n\
 Arguments: --output <filename> [--verbose] [--no-reboot] [--stdout-errors]\n\
-    [--delay <ms>]\n\
+    [--delay <ms>] [--usb-log-level <none/error/warning/debug>]\n\
 Description: Downloads the connected device's PIT file to the specified\n\
     output file.\n";
 
@@ -43,9 +44,11 @@ int DownloadPitAction::Execute(int argc, char **argv)
 	map<string, ArgumentType> argumentTypes;
 	argumentTypes["output"] = kArgumentTypeString;
 	argumentTypes["no-reboot"] = kArgumentTypeFlag;
+	argumentTypes["resume"] = kArgumentTypeFlag;
 	argumentTypes["delay"] = kArgumentTypeUnsignedInteger;
 	argumentTypes["verbose"] = kArgumentTypeFlag;
 	argumentTypes["stdout-errors"] = kArgumentTypeFlag;
+	argumentTypes["usb-log-level"] = kArgumentTypeString;
 
 	Arguments arguments(argumentTypes);
 
@@ -67,11 +70,48 @@ int DownloadPitAction::Execute(int argc, char **argv)
 	const UnsignedIntegerArgument *communicationDelayArgument = static_cast<const UnsignedIntegerArgument *>(arguments.GetArgument("delay"));
 
 	bool reboot = arguments.GetArgument("no-reboot") == nullptr;
+	bool resume = arguments.GetArgument("resume") != nullptr;
 	bool verbose = arguments.GetArgument("verbose") != nullptr;
 	
 	if (arguments.GetArgument("stdout-errors") != nullptr)
 		Interface::SetStdoutErrors(true);
 
+	const StringArgument *usbLogLevelArgument = static_cast<const StringArgument *>(arguments.GetArgument("usb-log-level"));
+
+	BridgeManager::UsbLogLevel usbLogLevel = BridgeManager::UsbLogLevel::Default;
+
+	if (usbLogLevelArgument)
+	{
+		const string& usbLogLevelString = usbLogLevelArgument->GetValue();
+
+		if (usbLogLevelString.compare("none") == 0 || usbLogLevelString.compare("NONE") == 0)
+		{
+			usbLogLevel = BridgeManager::UsbLogLevel::None;
+		}
+		else if (usbLogLevelString.compare("error") == 0 || usbLogLevelString.compare("ERROR") == 0)
+		{
+			usbLogLevel = BridgeManager::UsbLogLevel::Error;
+		}
+		else if (usbLogLevelString.compare("warning") == 0 || usbLogLevelString.compare("WARNING") == 0)
+		{
+			usbLogLevel = BridgeManager::UsbLogLevel::Warning;
+		}
+		else if (usbLogLevelString.compare("info") == 0 || usbLogLevelString.compare("INFO") == 0)
+		{
+			usbLogLevel = BridgeManager::UsbLogLevel::Info;
+		}
+		else if (usbLogLevelString.compare("debug") == 0 || usbLogLevelString.compare("DEBUG") == 0)
+		{
+			usbLogLevel = BridgeManager::UsbLogLevel::Debug;
+		}
+		else
+		{
+			Interface::Print("Unknown USB log level: %s\n\n", usbLogLevelString.c_str());
+			Interface::Print(DownloadPitAction::usage);
+			return (0);
+		}
+	}
+
 	// Info
 
 	Interface::PrintReleaseInfo();
@@ -96,8 +136,9 @@ int DownloadPitAction::Execute(int argc, char **argv)
 		communicationDelay = communicationDelayArgument->GetValue();
 
 	BridgeManager *bridgeManager = new BridgeManager(verbose, communicationDelay);
+	bridgeManager->SetUsbLogLevel(usbLogLevel);
 
-	if (bridgeManager->Initialise() != BridgeManager::kInitialiseSucceeded || !bridgeManager->BeginSession())
+	if (bridgeManager->Initialise(resume) != BridgeManager::kInitialiseSucceeded || !bridgeManager->BeginSession())
 	{
 		fclose(outputPitFile);
 		delete bridgeManager;