15 files changed, 451 insertions, 315 deletions

diff --git a/heimdall-frontend/Source/mainwindow.cpp b/heimdall-frontend/Source/mainwindow.cpp
index f3732c4..14ae787 100644
--- a/heimdall-frontend/Source/mainwindow.cpp
+++ b/heimdall-frontend/Source/mainwindow.cpp
@@ -1022,6 +1022,8 @@ void MainWindow::DownloadPit(void)
 	arguments.append("--output");

 	arguments.append(pitDestinationLineEdit->text());

 

+	arguments.append("--no-reboot");

+

 	arguments.append("--stdout-errors");

 

 	StartHeimdall(arguments);

diff --git a/heimdall-frontend/aboutform.ui b/heimdall-frontend/aboutform.ui
index b7a85e3..e0e6994 100644
--- a/heimdall-frontend/aboutform.ui
+++ b/heimdall-frontend/aboutform.ui
@@ -72,9 +72,9 @@
     <property name="geometry">

      <rect>

       <x>0</x>

-      <y>-8</y>

-      <width>544</width>

-      <height>776</height>

+      <y>0</y>

+      <width>542</width>

+      <height>1140</height>

      </rect>

     </property>

     <property name="sizePolicy">

@@ -112,14 +112,14 @@
           <string>&lt;!DOCTYPE HTML PUBLIC &quot;-//W3C//DTD HTML 4.0//EN&quot; &quot;http://www.w3.org/TR/REC-html40/strict.dtd&quot;&gt;

 &lt;html&gt;&lt;head&gt;&lt;meta name=&quot;qrichtext&quot; content=&quot;1&quot; /&gt;&lt;style type=&quot;text/css&quot;&gt;

 p, li { white-space: pre-wrap; }

-&lt;/style&gt;&lt;/head&gt;&lt;body style=&quot; font-family:'Lucida Grande'; font-size:13pt; font-weight:400; font-style:normal;&quot;&gt;

-&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-family:'MS Shell Dlg 2'; font-size:10pt; font-weight:600;&quot;&gt;Heimdall Frontend&lt;/span&gt;&lt;/p&gt;

-&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-family:'MS Shell Dlg 2'; font-size:10pt;&quot;&gt;Version 1.3.0&lt;/span&gt;&lt;/p&gt;

-&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-family:'MS Shell Dlg 2'; font-size:10pt;&quot;&gt;Copyright © 2010-2011 Benjamin Dobell, Glass Echidna&lt;/span&gt;&lt;/p&gt;

-&lt;p style=&quot;-qt-paragraph-type:empty; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px; font-family:'MS Shell Dlg 2'; font-size:10pt;&quot;&gt;&lt;/p&gt;

-&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-family:'MS Shell Dlg 2'; font-size:10pt; font-weight:600;&quot;&gt;Heimdall (command line)&lt;/span&gt;&lt;/p&gt;

-&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-family:'MS Shell Dlg 2'; font-size:10pt;&quot;&gt;Copyright © 2010-2011 Benjamin Dobell, Glass Echidna&lt;/span&gt;&lt;/p&gt;

-&lt;p style=&quot;-qt-paragraph-type:empty; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px; font-family:'MS Shell Dlg 2'; font-size:10pt;&quot;&gt;&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>

+&lt;/style&gt;&lt;/head&gt;&lt;body style=&quot; font-family:'MS Shell Dlg 2'; font-size:10pt; font-weight:400; font-style:normal;&quot;&gt;

+&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-weight:600;&quot;&gt;Heimdall Frontend&lt;/span&gt;&lt;/p&gt;

+&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;Version 1.3.1&lt;/p&gt;

+&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;Copyright © 2010-2011 Benjamin Dobell, Glass Echidna&lt;/p&gt;

+&lt;p style=&quot;-qt-paragraph-type:empty; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;/p&gt;

+&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;span style=&quot; font-weight:600;&quot;&gt;Heimdall (command line)&lt;/span&gt;&lt;/p&gt;

+&lt;p style=&quot; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;Copyright © 2010-2011 Benjamin Dobell, Glass Echidna&lt;/p&gt;

+&lt;p style=&quot;-qt-paragraph-type:empty; margin-top:0px; margin-bottom:0px; margin-left:0px; margin-right:0px; -qt-block-indent:0; text-indent:0px;&quot;&gt;&lt;/p&gt;&lt;/body&gt;&lt;/html&gt;</string>

          </property>

          <property name="textFormat">

           <enum>Qt::RichText</enum>

diff --git a/heimdall-frontend/mainwindow.ui b/heimdall-frontend/mainwindow.ui
index 541b610..987a99e 100644
--- a/heimdall-frontend/mainwindow.ui
+++ b/heimdall-frontend/mainwindow.ui
@@ -1718,7 +1718,7 @@
      <x>0</x>

      <y>0</y>

      <width>788</width>

-     <height>22</height>

+     <height>21</height>

     </rect>

    </property>

    <widget class="QMenu" name="menuHelp">

diff --git a/heimdall/heimdall.vcxproj b/heimdall/heimdall.vcxproj
index f440acb..47eb2db 100644
--- a/heimdall/heimdall.vcxproj
+++ b/heimdall/heimdall.vcxproj
@@ -103,14 +103,15 @@
     <ClInclude Include="source\BeginDumpPacket.h" />

     <ClInclude Include="source\BridgeManager.h" />

     <ClInclude Include="source\ControlPacket.h" />

-    <ClInclude Include="source\DeviceInfoPacket.h" />

-    <ClInclude Include="source\DeviceInfoResponse.h" />

+    <ClInclude Include="source\SetupSessionPacket.h" />

+    <ClInclude Include="source\SetupSessionResponse.h" />

     <ClInclude Include="source\DumpPartFileTransferPacket.h" />

     <ClInclude Include="source\DumpPartPitFilePacket.h" />

     <ClInclude Include="source\DumpResponse.h" />

     <ClInclude Include="source\EndFileTransferPacket.h" />

     <ClInclude Include="source\EndModemFileTransferPacket.h" />

     <ClInclude Include="source\EndPhoneFileTransferPacket.h" />

+    <ClInclude Include="source\EndPitFileTransferPacket.h" />

     <ClInclude Include="source\FileTransferPacket.h" />

     <ClInclude Include="source\FlashPartFileTransferPacket.h" />

     <ClInclude Include="source\FlashPartPitFilePacket.h" />

diff --git a/heimdall/heimdall.vcxproj.filters b/heimdall/heimdall.vcxproj.filters
index e89df24..3d2faa0 100644
--- a/heimdall/heimdall.vcxproj.filters
+++ b/heimdall/heimdall.vcxproj.filters
@@ -16,12 +16,6 @@
     <ClInclude Include="source\ControlPacket.h">

       <Filter>Source</Filter>

     </ClInclude>

-    <ClInclude Include="source\DeviceInfoPacket.h">

-      <Filter>Source</Filter>

-    </ClInclude>

-    <ClInclude Include="source\DeviceInfoResponse.h">

-      <Filter>Source</Filter>

-    </ClInclude>

     <ClInclude Include="source\DumpPartFileTransferPacket.h">

       <Filter>Source</Filter>

     </ClInclude>

@@ -79,10 +73,19 @@
     <ClInclude Include="source\SendFilePartResponse.h">

       <Filter>Source</Filter>

     </ClInclude>

-    <ClInclude Include="source\EndSessionPacket.h">

+    <ClInclude Include="source\Interface.h">

       <Filter>Source</Filter>

     </ClInclude>

-    <ClInclude Include="source\Interface.h">

+    <ClInclude Include="source\EndPitFileTransferPacket.h">

+      <Filter>Source</Filter>

+    </ClInclude>

+    <ClInclude Include="source\SetupSessionPacket.h">

+      <Filter>Source</Filter>

+    </ClInclude>

+    <ClInclude Include="source\SetupSessionResponse.h">

+      <Filter>Source</Filter>

+    </ClInclude>

+    <ClInclude Include="source\EndSessionPacket.h">

       <Filter>Source</Filter>

     </ClInclude>

   </ItemGroup>

diff --git a/heimdall/source/BridgeManager.cpp b/heimdall/source/BridgeManager.cpp
index 00bdf29..db522ec 100644
--- a/heimdall/source/BridgeManager.cpp
+++ b/heimdall/source/BridgeManager.cpp
@@ -27,13 +27,14 @@
 // Heimdall
 #include "BeginDumpPacket.h"
 #include "BridgeManager.h"
-#include "DeviceInfoPacket.h"
-#include "DeviceInfoResponse.h"
+#include "SetupSessionPacket.h"
+#include "SetupSessionResponse.h"
 #include "DumpPartFileTransferPacket.h"
 #include "DumpPartPitFilePacket.h"
 #include "DumpResponse.h"
 #include "EndModemFileTransferPacket.h"
 #include "EndPhoneFileTransferPacket.h"
+#include "EndPitFileTransferPacket.h"
 #include "EndSessionPacket.h"
 #include "FileTransferPacket.h"
 #include "FlashPartFileTransferPacket.h"
@@ -66,6 +67,159 @@ enum
 	kMaxSequenceLength = 800
 };
 
+bool BridgeManager::CheckProtocol(void) const
+{
+	Interface::Print("Checking if protocol is initialised...\n");
+
+	SetupSessionPacket deviceInfoPacket(SetupSessionPacket::kDeviceInfo);
+
+	if (!SendPacket(&deviceInfoPacket, 100, false))
+	{
+		Interface::Print("Protocol is not initialised.\n");
+		return (false);
+	}
+
+	SetupSessionResponse deviceInfoResponse;
+
+	if (!ReceivePacket(&deviceInfoResponse, 100, false))
+	{
+		Interface::Print("Protocol is not initialised.\n");
+		return (false);
+	}
+
+	Interface::Print("Protocol is initialised.\n");
+	return (true);
+}
+
+bool BridgeManager::InitialiseProtocol(void) const
+{
+	Interface::Print("Initialising protocol...\n");
+
+	unsigned char *dataBuffer = new unsigned char[7];
+
+	int result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x3, 0, nullptr, 0, 1000);
+
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to initialise protocol!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	memset(dataBuffer, 0, 7);
+	dataBuffer[1] = 0xC2;
+	dataBuffer[2] = 0x01;
+	dataBuffer[6] = 0x07;
+
+	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x20, 0x0, 0, dataBuffer, 7, 1000);
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to initialise protocol!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x3, 0, nullptr, 0, 1000);
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to initialise protocol!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x2, 0, nullptr, 0, 1000);
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to initialise protocol!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	memset(dataBuffer, 0, 7);
+	dataBuffer[1] = 0xC2;
+	dataBuffer[2] = 0x01;
+	dataBuffer[6] = 0x08;
+
+	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x20, 0x0, 0, dataBuffer, 7, 1000);
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to initialise protocol!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x2, 0, nullptr, 0, 1000);
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to initialise protocol!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	Interface::Print("Handshaking with Loke...\n");
+
+	int dataTransferred;
+
+	// Send "ODIN"
+	strcpy((char *)dataBuffer, "ODIN");
+
+	result = libusb_bulk_transfer(deviceHandle, outEndpoint, dataBuffer, 4, &dataTransferred, 1000);
+	if (result < 0)
+	{
+		if (verbose)
+			Interface::PrintError("Failed to send data: \"%s\"\n", dataBuffer);
+		else
+			Interface::PrintError("Failed to send data!");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	if (dataTransferred != 4)
+	{
+		if (verbose)
+			Interface::PrintError("Failed to complete sending of data: \"%s\"\n", dataBuffer);
+		else
+			Interface::PrintError("Failed to complete sending of data!");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	// Expect "LOKE"
+	memset(dataBuffer, 0, 7);
+
+	result = libusb_bulk_transfer(deviceHandle, inEndpoint, dataBuffer, 7, &dataTransferred, 1000);
+	if (result < 0)
+	{
+		Interface::PrintError("Failed to receive response!\n");
+
+		delete [] dataBuffer;
+		return (false);;
+	}
+
+	if (dataTransferred != 4 || memcmp(dataBuffer, "LOKE", 4) != 0)
+	{
+		Interface::PrintError("Unexpected communication!\n");
+
+		if (verbose)
+			Interface::PrintError("Expected: \"%s\"\nReceived: \"%s\"\n", "LOKE", dataBuffer);
+
+		Interface::PrintError("Handshake failed!\n");
+
+		delete [] dataBuffer;
+		return (false);
+	}
+
+	return (true);
+}
+
 BridgeManager::BridgeManager(bool verbose, int communicationDelay)
 {
 	this->verbose = verbose;
@@ -339,6 +493,14 @@ int BridgeManager::Initialise(void)
 
 	Interface::Print("\n");
 
+	if (!CheckProtocol())
+	{
+		if (!InitialiseProtocol())
+			return (BridgeManager::kInitialiseFailed);
+	}
+
+	Interface::Print("\n");
+
 	return (BridgeManager::kInitialiseSucceeded);
 }
 
@@ -346,130 +508,53 @@ bool BridgeManager::BeginSession(void) const
 {
 	Interface::Print("Beginning session...\n");
 
-	unsigned char *dataBuffer = new unsigned char[7];
-
-	int result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x3, 0, nullptr, 0, 1000);
+	SetupSessionPacket beginSessionPacket(SetupSessionPacket::kBeginSession);
 
-	if (result < 0)
+	if (!SendPacket(&beginSessionPacket))
 	{
 		Interface::PrintError("Failed to begin session!\n");
-
-		delete [] dataBuffer;
 		return (false);
 	}
 
-	memset(dataBuffer, 0, 7);
-	dataBuffer[1] = 0xC2;
-	dataBuffer[2] = 0x01;
-	dataBuffer[6] = 0x07;
-
-	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x20, 0x0, 0, dataBuffer, 7, 1000);
-	if (result < 0)
-	{
-		Interface::PrintError("Failed to begin session!\n");
-
-		delete [] dataBuffer;
+	SetupSessionResponse setupSessionResponse;
+	if (!ReceivePacket(&setupSessionResponse))
 		return (false);
-	}
 
-	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x3, 0, nullptr, 0, 1000);
-	if (result < 0)
-	{
-		Interface::PrintError("Failed to begin session!\n");
+	int result = setupSessionResponse.GetUnknown();
 
-		delete [] dataBuffer;
-		return (false);
-	}
-
-	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x2, 0, nullptr, 0, 1000);
-	if (result < 0)
-	{
-		Interface::PrintError("Failed to begin session!\n");
-
-		delete [] dataBuffer;
-		return (false);
-	}
-
-	memset(dataBuffer, 0, 7);
-	dataBuffer[1] = 0xC2;
-	dataBuffer[2] = 0x01;
-	dataBuffer[6] = 0x08;
-
-	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x20, 0x0, 0, dataBuffer, 7, 1000);
-	if (result < 0)
-	{
-		Interface::PrintError("Failed to begin session!\n");
-
-		delete [] dataBuffer;
-		return (false);
-	}
-
-	result = libusb_control_transfer(deviceHandle, LIBUSB_REQUEST_TYPE_CLASS, 0x22, 0x2, 0, nullptr, 0, 1000);
-	if (result < 0)
+	// 131072 for Galaxy S II, 0 for other devices.
+	if (result != 0 && result != 131072)
 	{
-		Interface::PrintError("Failed to begin session!\n");
-
-		delete [] dataBuffer;
+		Interface::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%d\n", result);
 		return (false);
 	}
 
-	Interface::Print("Handshaking with Loke...\n");
-
-	int dataTransferred;
+	// -------------------- KIES DOESN'T DO THIS --------------------
 
-	// Send "ODIN"
-	strcpy((char *)dataBuffer, "ODIN");
+	SetupSessionPacket deviceTypePacket(SetupSessionPacket::kDeviceInfo);
 
-	result = libusb_bulk_transfer(deviceHandle, outEndpoint, dataBuffer, 4, &dataTransferred, 1000);
-	if (result < 0)
+	if (!SendPacket(&deviceTypePacket))
 	{
-		if (verbose)
-			Interface::PrintError("Failed to send data: \"%s\"\n", dataBuffer);
-		else
-			Interface::PrintError("Failed to send data!");
-
-		delete [] dataBuffer;
+		Interface::PrintError("Failed to request device type!\n");
 		return (false);
 	}
 
-	if (dataTransferred != 4)
-	{
-		if (verbose)
-			Interface::PrintError("Failed to complete sending of data: \"%s\"\n", dataBuffer);
-		else
-			Interface::PrintError("Failed to complete sending of data!");
-
-		delete [] dataBuffer;
+	if (!ReceivePacket(&setupSessionResponse))
 		return (false);
-	}
 
-	// Expect "LOKE"
-	memset(dataBuffer, 0, 7);
+	int deviceType = setupSessionResponse.GetUnknown();
 
-	result = libusb_bulk_transfer(deviceHandle, inEndpoint, dataBuffer, 7, &dataTransferred, 1000);
-	if (result < 0)
+	// 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 (deviceType != 180 && deviceType != 0 && deviceType != 3 && deviceType != 190)
 	{
-		Interface::PrintError("Failed to receive response!\n");
-
-		delete [] dataBuffer;
-		return (false);;
+		Interface::PrintError("Unexpected device info response!\nExpected: 180, 0 or 3\nReceived:%d\n", deviceType);
+		return (false);
 	}
-
-	if (dataTransferred != 4 || memcmp(dataBuffer, "LOKE", 4) != 0)
+	else
 	{
-		Interface::PrintError("Unexpected communication!\n");
-
-		if (verbose)
-			Interface::PrintError("Expected: \"%s\"\nReceived: \"%s\"\n", "LOKE", dataBuffer);
-
-		Interface::PrintError("Handshake failed!\n");
-
-		delete [] dataBuffer;
-		return (false);
+		Interface::Print("Session begun with device of type: %d\n\n", result);
 	}
 
-	Interface::Print("\n");
-
 	return (true);
 }
 
@@ -529,7 +614,7 @@ bool BridgeManager::EndSession(bool reboot) const
 	return (true);
 }
 
-bool BridgeManager::SendPacket(OutboundPacket *packet, int timeout) const
+bool BridgeManager::SendPacket(OutboundPacket *packet, int timeout, bool retry) const
 {
 	packet->Pack();
 
@@ -537,7 +622,7 @@ bool BridgeManager::SendPacket(OutboundPacket *packet, int timeout) const
 	int result = libusb_bulk_transfer(deviceHandle, outEndpoint, packet->GetData(), packet->GetSize(),
 		&dataTransferred, timeout);
 
-	if (result < 0)
+	if (result < 0 && retry)
 	{
 		// max(250, communicationDelay)
 		int retryDelay = (communicationDelay > 250) ? communicationDelay : 250;
@@ -577,13 +662,13 @@ bool BridgeManager::SendPacket(OutboundPacket *packet, int timeout) const
 	return (true);
 }
 
-bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout) const
+bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout, bool retry) const
 {
 	int dataTransferred;
 	int result = libusb_bulk_transfer(deviceHandle, inEndpoint, packet->GetData(), packet->GetSize(),
 		&dataTransferred, timeout);
 
-	if (result < 0)
+	if (result < 0 && retry)
 	{
 		// max(250, communicationDelay)
 		int retryDelay = (communicationDelay > 250) ? communicationDelay : 250;
@@ -608,12 +693,6 @@ bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout) const
 
 			if (verbose)
 				Interface::PrintError("libusb error %d whilst receiving packet.", result);
-
-			if (i >= 3)
-			{
-				int breakHere = 0;
-				breakHere++;
-			}
 		}
 
 		if (verbose)
@@ -633,8 +712,8 @@ bool BridgeManager::ReceivePacket(InboundPacket *packet, int timeout) const
 
 bool BridgeManager::RequestDeviceInfo(unsigned int request, int *result) const
 {
-	DeviceInfoPacket deviceInfoPacket(request);
-	bool success = SendPacket(&deviceInfoPacket);
+	SetupSessionPacket beginSessionPacket(request);
+	bool success = SendPacket(&beginSessionPacket);
 
 	if (!success)
 	{
@@ -646,8 +725,10 @@ bool BridgeManager::RequestDeviceInfo(unsigned int request, int *result) const
 		return (false);
 	}
 
-	DeviceInfoResponse deviceInfoResponse;
-	success = ReceivePacket(&deviceInfoResponse);
+	SetupSessionResponse deviceInfoResponse;
+	if (!ReceivePacket(&deviceInfoResponse))
+		return (false);
+
 	*result = deviceInfoResponse.GetUnknown();
 
 	return (true);
@@ -718,7 +799,28 @@ bool BridgeManager::SendPitFile(FILE *file) const
 
 	if (!success)
 	{
-		Interface::PrintError("Failed to receive PIT file transfer count!\n");
+		Interface::PrintError("Failed to receive PIT file part response!\n");
+		return (false);
+	}
+
+	// End pit file transfer
+	EndPitFileTransferPacket *endPitFileTransferPacket = new EndPitFileTransferPacket(fileSize);
+	success = SendPacket(endPitFileTransferPacket);
+	delete endPitFileTransferPacket;
+
+	if (!success)
+	{
+		Interface::PrintError("Failed to send end PIT file transfer packet!\n");
+		return (false);
+	}
+
+	pitFileResponse = new PitFileResponse();
+	success = ReceivePacket(pitFileResponse);
+	delete pitFileResponse;
+
+	if (!success)
+	{
+		Interface::PrintError("Failed to confirm end of PIT file transfer!\n");
 		return (false);
 	}
 
diff --git a/heimdall/source/BridgeManager.h b/heimdall/source/BridgeManager.h
index d87f67d..f929ac9 100644
--- a/heimdall/source/BridgeManager.h
+++ b/heimdall/source/BridgeManager.h
@@ -99,6 +99,9 @@ namespace Heimdall
 

 #endif

 

+			bool CheckProtocol(void) const;

+			bool InitialiseProtocol(void) const;

+

 		public:

 

 			BridgeManager(bool verbose, int communicationDelay);

@@ -110,8 +113,8 @@ namespace Heimdall
 			bool BeginSession(void) const;

 			bool EndSession(bool reboot) const;

 

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

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

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

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

 

 			bool RequestDeviceInfo(unsigned int request, int *result) const;

 

diff --git a/heimdall/source/ControlPacket.h b/heimdall/source/ControlPacket.h
index f003654..ce957b2 100644
--- a/heimdall/source/ControlPacket.h
+++ b/heimdall/source/ControlPacket.h
@@ -32,7 +32,7 @@ namespace Heimdall
 

 			enum

 			{

-				kControlTypeDeviceInfo      = 0x64,

+				kControlTypeSetupSession      = 0x64,

 				kControlTypePitFile         = 0x65,

 				kControlTypeFileTransfer    = 0x66,

 				kControlTypeEndSession      = 0x67

diff --git a/heimdall/source/EndPitFileTransferPacket.h b/heimdall/source/EndPitFileTransferPacket.h
new file mode 100644
index 0000000..38b7317
--- /dev/null
+++ b/heimdall/source/EndPitFileTransferPacket.h
@@ -0,0 +1,56 @@
+/* Copyright (c) 2010-2011 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 ENDPITFILETRANSFERPACKET_H
+#define ENDPITFILETRANSFERPACKET_H
+
+// Heimdall
+#include "PitFilePacket.h"
+
+namespace Heimdall
+{
+	class EndPitFileTransferPacket : public PitFilePacket
+	{
+		private:
+
+			unsigned int fileSize;
+
+		public:
+
+			EndPitFileTransferPacket(unsigned int fileSize) : PitFilePacket(PitFilePacket::kRequestEndTransfer)
+			{
+				this->fileSize = fileSize;
+			}
+
+			unsigned int GetFileSize(void) const
+			{
+				return (fileSize);
+			}
+
+			void Pack(void)
+			{
+				PitFilePacket::Pack();
+
+				PackInteger(PitFilePacket::kDataSize, fileSize);
+			}
+	};
+}
+
+#endif
diff --git a/heimdall/source/EndSessionPacket.h b/heimdall/source/EndSessionPacket.h
index 6547efb..58b90e7 100644
--- a/heimdall/source/EndSessionPacket.h
+++ b/heimdall/source/EndSessionPacket.h
@@ -18,8 +18,8 @@
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  THE SOFTWARE.*/
 
-#ifndef REBOOTDEVICEPACKET_H
-#define REBOOTDEVICEPACKET_H
+#ifndef ENDSESSIONPACKET_H
+#define ENDSESSIONPACKET_H
 
 // Heimdall
 #include "ControlPacket.h"
diff --git a/heimdall/source/Interface.cpp b/heimdall/source/Interface.cpp
index 8f2d313..55aa392 100644
--- a/heimdall/source/Interface.cpp
+++ b/heimdall/source/Interface.cpp
@@ -33,7 +33,7 @@ using namespace Heimdall;
 
 bool Interface::stdoutErrors = false;
 
-const char *Interface::version = "v1.3.0";
+const char *Interface::version = "v1.3.1";
 
 const char *Interface::usage = "Usage: heimdall <action> <action arguments> <common arguments>\n\
 \n\
diff --git a/heimdall/source/ResponsePacket.h b/heimdall/source/ResponsePacket.h
index 6094c83..230b28c 100644
--- a/heimdall/source/ResponsePacket.h
+++ b/heimdall/source/ResponsePacket.h
@@ -33,7 +33,7 @@ namespace Heimdall
 			enum

 			{

 				kResponseTypeSendFilePart   = 0x00,

-				kResponseTypeDeviceInfo     = 0x64,

+				kResponseTypeBeginSession     = 0x64,

 				kResponseTypePitFile        = 0x65,

 				kResponseTypeFileTransfer   = 0x66,

 				kResponseTypeEndSession     = 0x67

diff --git a/heimdall/source/DeviceInfoPacket.h b/heimdall/source/SetupSessionPacket.h
index 9ec1be3..cb75a35 100644
--- a/heimdall/source/DeviceInfoPacket.h
+++ b/heimdall/source/SetupSessionPacket.h
@@ -1,74 +1,74 @@
-/* Copyright (c) 2010-2011 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 DEVICEINFOPACKET_H

-#define DEVICEINFOPACKET_H

-

-// Heimdall

-#include "ControlPacket.h"

-

-namespace Heimdall

-{

-	class DeviceInfoPacket : public ControlPacket

-	{

-		public:

-

-			enum

-			{

-				kUnknown1			= 0,	// Begin device info?

-				kUnknown2			= 1,	// End device info?

-				kTotalBytes			= 2		// Total no. bytes that will be flashed, seems to be ignored though.

-			};

-

-		private:

-

-			unsigned int request;

-			unsigned int unknown3Parameter;	// TODO: Subclass for unknown2.

-

-		public:

-

-			DeviceInfoPacket(unsigned int request, unsigned int unknown3Parameter = 0)

-				: ControlPacket(ControlPacket::kControlTypeDeviceInfo)

-			{

-				this->request = request;

-				this->unknown3Parameter = unknown3Parameter;

-			}

-

-			unsigned int GetRequest(void) const

-			{

-				return (request);

-			}

-

-			unsigned int GetUnknown3Parameter(void) const

-			{

-				return (unknown3Parameter);

-			}

-

-			void Pack(void)

-			{

-				ControlPacket::Pack();

-

-				PackInteger(ControlPacket::kDataSize, request);

-				PackInteger(ControlPacket::kDataSize + 4, unknown3Parameter);

-			}

-	};

-}

-

-#endif

+/* Copyright (c) 2010-2011 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 SETUPSESSIONPACKET_H
+#define SETUPSESSIONPACKET_H
+
+// Heimdall
+#include "ControlPacket.h"
+
+namespace Heimdall
+{
+	class SetupSessionPacket : public ControlPacket
+	{
+		public:
+
+			enum
+			{
+				kBeginSession		= 0,
+				kDeviceInfo			= 1,
+				kTotalBytes			= 2
+			};
+
+		private:
+
+			unsigned int request;
+			unsigned int unknown3Parameter;
+
+		public:
+
+			SetupSessionPacket(unsigned int request, unsigned int unknown3Parameter = 0)
+				: ControlPacket(ControlPacket::kControlTypeSetupSession)
+			{
+				this->request = request;
+				this->unknown3Parameter = unknown3Parameter;
+			}
+
+			unsigned int GetRequest(void) const
+			{
+				return (request);
+			}
+
+			unsigned int GetUnknown3Parameter(void) const
+			{
+				return (unknown3Parameter);
+			}
+
+			void Pack(void)
+			{
+				ControlPacket::Pack();
+
+				PackInteger(ControlPacket::kDataSize, request);
+				PackInteger(ControlPacket::kDataSize + 4, unknown3Parameter);
+			}
+	};
+}
+
+#endif
diff --git a/heimdall/source/DeviceInfoResponse.h b/heimdall/source/SetupSessionResponse.h
index 11a0404..a107f29 100644
--- a/heimdall/source/DeviceInfoResponse.h
+++ b/heimdall/source/SetupSessionResponse.h
@@ -1,58 +1,58 @@
-/* Copyright (c) 2010-2011 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 DEVICEINFORESPONSE_H

-#define DEVICEINFORESPONSE_H

-

-// Heimdall

-#include "ResponsePacket.h"

-

-namespace Heimdall

-{

-	class DeviceInfoResponse : public ResponsePacket

-	{

-		private:

-

-			unsigned int unknown;

-

-		public:

-

-			DeviceInfoResponse() : ResponsePacket(ResponsePacket::kResponseTypeDeviceInfo)

-			{

-			}

-

-			int GetUnknown(void) const

-			{

-				return (unknown);

-			}

-

-			bool Unpack(void)

-			{

-				if (!ResponsePacket::Unpack())

-					return (false);

-

-				unknown = UnpackInteger(ResponsePacket::kDataSize);

-

-				return (true);

-			}

-	};

-}

-

-#endif

+/* Copyright (c) 2010-2011 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 SETUPSESSIONRESPONSE_H
+#define SETUPSESSIONRESPONSE_H
+
+// Heimdall
+#include "ResponsePacket.h"
+
+namespace Heimdall
+{
+	class SetupSessionResponse : public ResponsePacket
+	{
+		private:
+
+			unsigned int unknown;
+
+		public:
+
+			SetupSessionResponse() : ResponsePacket(ResponsePacket::kResponseTypeBeginSession)
+			{
+			}
+
+			int GetUnknown(void) const
+			{
+				return (unknown);
+			}
+
+			bool Unpack(void)
+			{
+				if (!ResponsePacket::Unpack())
+					return (false);
+
+				unknown = UnpackInteger(ResponsePacket::kDataSize);
+
+				return (true);
+			}
+	};
+}
+
+#endif
diff --git a/heimdall/source/main.cpp b/heimdall/source/main.cpp
index a3ae962..2d05e33 100644
--- a/heimdall/source/main.cpp
+++ b/heimdall/source/main.cpp
@@ -29,8 +29,8 @@
 
 // Heimdall
 #include "BridgeManager.h"
-#include "DeviceInfoPacket.h"
-#include "DeviceInfoResponse.h"
+#include "SetupSessionPacket.h"
+#include "SetupSessionResponse.h"
 #include "EndModemFileTransferPacket.h"
 #include "EndPhoneFileTransferPacket.h"
 #include "Interface.h"
@@ -244,37 +244,6 @@ void closeFiles(map<string, FILE *> argumentfileMap)
 	argumentfileMap.clear();
 }
 
-bool retrieveDeviceInfo(BridgeManager *bridgeManager)
-{
-	// ---------- GET DEVICE INFORMATION ----------
-
-	int deviceInfoResult;
-	
-	if (!bridgeManager->RequestDeviceInfo(DeviceInfoPacket::kUnknown1, &deviceInfoResult))
-		return (false);
-
-	// 131072 for Galaxy S II, 0 for other devices.
-	if (deviceInfoResult != 0 && deviceInfoResult != 131072)
-	{
-		Interface::PrintError("Unexpected device info response!\nExpected: 0\nReceived:%d\n", deviceInfoResult);
-		return (false);
-	}
-
-	// -------------------- KIES DOESN'T DO THIS --------------------
-
-	if (!bridgeManager->RequestDeviceInfo(DeviceInfoPacket::kUnknown2, &deviceInfoResult))
-		return (false);
-
-	// 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:%d\n", deviceInfoResult);
-		return (false);
-	}
-
-	return (true);
-}
-
 int downloadPitFile(BridgeManager *bridgeManager, unsigned char **pitBuffer)
 {
 	Interface::Print("Downloading device's PIT file...\n");
@@ -390,7 +359,7 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 		}
 	}
 	
-	DeviceInfoPacket *deviceInfoPacket = new DeviceInfoPacket(DeviceInfoPacket::kTotalBytes, totalBytes);
+	SetupSessionPacket *deviceInfoPacket = new SetupSessionPacket(SetupSessionPacket::kTotalBytes, totalBytes);
 	success = bridgeManager->SendPacket(deviceInfoPacket);
 	delete deviceInfoPacket;
 
@@ -400,7 +369,7 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 		return (false);
 	}
 
-	DeviceInfoResponse *deviceInfoResponse = new DeviceInfoResponse();
+	SetupSessionResponse *deviceInfoResponse = new SetupSessionResponse();
 	success = bridgeManager->ReceivePacket(deviceInfoResponse);
 	int deviceInfoResult = deviceInfoResponse->GetUnknown();
 	delete deviceInfoResponse;
@@ -670,7 +639,7 @@ int main(int argc, char **argv)
 				return (0);
 			}
 
-			if (!bridgeManager->BeginSession() || !retrieveDeviceInfo(bridgeManager))
+			if (!bridgeManager->BeginSession())
 			{
 				closeFiles(argumentFileMap);
 				delete bridgeManager;
@@ -690,7 +659,7 @@ int main(int argc, char **argv)
 
 		case Interface::kActionClosePcScreen:
 		{
-			if (!bridgeManager->BeginSession() || !retrieveDeviceInfo(bridgeManager))
+			if (!bridgeManager->BeginSession())
 			{
 				delete bridgeManager;
 				return (-1);
@@ -717,7 +686,7 @@ int main(int argc, char **argv)
 				return (0);
 			}
 
-			if (!bridgeManager->BeginSession() || !retrieveDeviceInfo(bridgeManager))
+			if (!bridgeManager->BeginSession())
 			{
 				delete bridgeManager;
 				fclose(outputPitFile);
@@ -768,7 +737,7 @@ int main(int argc, char **argv)
 
 			int chipId = atoi(argumentMap.find(Interface::actions[Interface::kActionDump].valueArguments[Interface::kDumpValueArgChipId])->second.c_str());
 
-			if (!bridgeManager->BeginSession() || !retrieveDeviceInfo(bridgeManager))
+			if (!bridgeManager->BeginSession())
 			{
 				fclose(dumpFile);
 
@@ -787,7 +756,7 @@ int main(int argc, char **argv)
 
 		case Interface::kActionPrintPit:
 		{
-			if (!bridgeManager->BeginSession() || !retrieveDeviceInfo(bridgeManager))
+			if (!bridgeManager->BeginSession())
 			{
 				delete bridgeManager;
 				return (-1);