8 files changed, 393 insertions, 100 deletions

diff --git a/heimdall-frontend/Source/Alerts.cpp b/heimdall-frontend/Source/Alerts.cpp
index e8e8752..5dc0337 100644
--- a/heimdall-frontend/Source/Alerts.cpp
+++ b/heimdall-frontend/Source/Alerts.cpp
@@ -29,6 +29,7 @@ using namespace HeimdallFrontend;
 void Alerts::DisplayError(const QString& errorMessage)
 {
 	QMessageBox messageBox;
+	messageBox.setModal(true);
 	messageBox.setText(errorMessage);
 	messageBox.setIcon(QMessageBox::Critical);
 	messageBox.exec();
@@ -37,6 +38,7 @@ void Alerts::DisplayError(const QString& errorMessage)
 void Alerts::DisplayWarning(const QString& warningMessage)
 {
 	QMessageBox messageBox;
+	messageBox.setModal(true);
 	messageBox.setText(warningMessage);
 	messageBox.setIcon(QMessageBox::Warning);
 	messageBox.exec();
diff --git a/heimdall-frontend/Source/FirmwareInfo.cpp b/heimdall-frontend/Source/FirmwareInfo.cpp
index aee9313..7870621 100644
--- a/heimdall-frontend/Source/FirmwareInfo.cpp
+++ b/heimdall-frontend/Source/FirmwareInfo.cpp
@@ -18,9 +18,13 @@
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  THE SOFTWARE.*/
 
+// Qt
+#include "QRegExp"
+
 // Heimdall Frontend
 #include "Alerts.h"
 #include "FirmwareInfo.h"
+#include "Packaging.h"
 
 using namespace HeimdallFrontend;
 
@@ -320,7 +324,7 @@ bool FileInfo::ParseXml(QXmlStreamReader& xml)
 	return (false);
 }
 
-void FileInfo::WriteXml(QXmlStreamWriter& xml) const
+void FileInfo::WriteXml(QXmlStreamWriter& xml, const QString& filename) const
 {
 	xml.writeStartElement("file");
 
@@ -329,14 +333,7 @@ void FileInfo::WriteXml(QXmlStreamWriter& xml) const
 	xml.writeEndElement();
 
 	xml.writeStartElement("filename");
-
-	int lastSlash = filename.lastIndexOf('/');
-
-	if (lastSlash < 0)
-		lastSlash = filename.lastIndexOf('\\');
-
-	xml.writeCharacters(filename.mid(lastSlash + 1));
-
+	xml.writeCharacters(filename);
 	xml.writeEndElement();
 
 	xml.writeEndElement();
@@ -775,7 +772,7 @@ void FirmwareInfo::WriteXml(QXmlStreamWriter& xml) const
 	xml.writeStartElement("files");
 
 	for (int i = 0; i < fileInfos.length(); i++)
-		fileInfos[i].WriteXml(xml);
+		fileInfos[i].WriteXml(xml, Packaging::ClashlessFilename(fileInfos, i));
 
 	xml.writeEndElement();
 
diff --git a/heimdall-frontend/Source/FirmwareInfo.h b/heimdall-frontend/Source/FirmwareInfo.h
index 64c73cb..5765199 100644
--- a/heimdall-frontend/Source/FirmwareInfo.h
+++ b/heimdall-frontend/Source/FirmwareInfo.h
@@ -126,7 +126,7 @@ namespace HeimdallFrontend
 			FileInfo(unsigned int partitionId, const QString& filename);
 
 			bool ParseXml(QXmlStreamReader& xml);
-			void WriteXml(QXmlStreamWriter& xml) const;
+			void WriteXml(QXmlStreamWriter& xml, const QString& filename) const;
 
 			unsigned int GetPartitionId(void) const
 			{
diff --git a/heimdall-frontend/Source/Packaging.cpp b/heimdall-frontend/Source/Packaging.cpp
index fc41ba6..6c3890b 100644
--- a/heimdall-frontend/Source/Packaging.cpp
+++ b/heimdall-frontend/Source/Packaging.cpp
@@ -221,12 +221,12 @@ bool Packaging::ExtractTar(QTemporaryFile& tarFile, PackageData *packageData)
 	return (true);
 }
 
-bool Packaging::WriteTarEntry(const QString& filename, QTemporaryFile *tarFile, bool firmwareXml)
+bool Packaging::WriteTarEntry(const QString& filePath, QTemporaryFile *tarFile, const QString& entryFilename)
 {
 	TarHeader tarHeader;
 	memset(tarHeader.buffer, 0, TarHeader::kBlockLength);
 
-	QFile file(filename);
+	QFile file(filePath);
 
 	if (!file.open(QFile::ReadOnly))
 	{
@@ -243,19 +243,12 @@ bool Packaging::WriteTarEntry(const QString& filename, QTemporaryFile *tarFile,
 	QFileInfo qtFileInfo(file);
 	QByteArray utfFilename;
 
-	if (firmwareXml)
-	{
-		utfFilename = QString("firmware.xml").toUtf8();
-	}
-	else
-	{
-		utfFilename = qtFileInfo.fileName().toUtf8();
+	utfFilename = entryFilename.toUtf8();
 
-		if (utfFilename.length() > 100)
-		{
-			Alerts::DisplayError(QString("File name is too long:\n%1").arg(qtFileInfo.fileName()));
-			return (false);
-		}
+	if (utfFilename.length() > 100)
+	{
+		Alerts::DisplayError(QString("File name is too long:\n%1").arg(qtFileInfo.fileName()));
+		return (false);
 	}
 
 	strcpy(tarHeader.fields.name, utfFilename.constData());
@@ -388,7 +381,15 @@ bool Packaging::CreateTar(const FirmwareInfo& firmwareInfo, QTemporaryFile *tarF
 
 	for (int i = 0; i < fileInfos.length(); i++)
 	{
-		if (!WriteTarEntry(fileInfos[i].GetFilename(), tarFile))
+		QString filename = ClashlessFilename(fileInfos, i);
+
+		if (filename == "firmware.xml")
+		{
+			Alerts::DisplayError("You cannot name your partition files \"firmware.xml\".\nIt is a reserved name.");
+			return (false);
+		}
+
+		if (!WriteTarEntry(fileInfos[i].GetFilename(), tarFile, ClashlessFilename(fileInfos, i)))
 		{
 			tarFile->resize(0);
 			tarFile->close();
@@ -411,7 +412,20 @@ bool Packaging::CreateTar(const FirmwareInfo& firmwareInfo, QTemporaryFile *tarF
 		}
 	}
 
-	if (!WriteTarEntry(firmwareInfo.GetPitFilename(), tarFile))
+	int lastSlash = firmwareInfo.GetPitFilename().lastIndexOf('/');
+
+	if (lastSlash < 0)
+		lastSlash = firmwareInfo.GetPitFilename().lastIndexOf('\\');
+
+	QString pitFilename = ClashlessFilename(fileInfos, firmwareInfo.GetPitFilename().mid(lastSlash + 1));
+
+	if (pitFilename == "firmware.xml")
+	{
+		Alerts::DisplayError("You cannot name your PIT file \"firmware.xml\".\nIt is a reserved name.");
+		return (false);
+	}
+
+	if (!WriteTarEntry(firmwareInfo.GetPitFilename(), tarFile, pitFilename))
 	{
 		tarFile->resize(0);
 		tarFile->close();
@@ -431,7 +445,7 @@ bool Packaging::CreateTar(const FirmwareInfo& firmwareInfo, QTemporaryFile *tarF
 		return (false);
 	}
 
-	if (!WriteTarEntry(firmwareXmlFile.fileName(), tarFile, true))
+	if (!WriteTarEntry(firmwareXmlFile.fileName(), tarFile, "firmware.xml"))
 	{
 		tarFile->resize(0);
 		tarFile->close();
@@ -558,7 +572,12 @@ bool Packaging::BuildPackage(const QString& packagePath, const FirmwareInfo& fir
 	QTemporaryFile tar("XXXXXX.tar");
 
 	if (!CreateTar(firmwareInfo, &tar))
+	{
+		fclose(compressedPackageFile);
+		remove(packagePath.toStdString().c_str());
+
 		return (false);
+	}
 
 	if (!tar.open())
 	{
@@ -626,3 +645,271 @@ bool Packaging::BuildPackage(const QString& packagePath, const FirmwareInfo& fir
 
 	return (true);
 }
+
+QString Packaging::ClashlessFilename(const QList<FileInfo>& fileInfos, int fileInfoIndex)
+{
+	int lastSlash = fileInfos[fileInfoIndex].GetFilename().lastIndexOf('/');
+
+	if (lastSlash < 0)
+		lastSlash = fileInfos[fileInfoIndex].GetFilename().lastIndexOf('\\');
+
+	QString filename = fileInfos[fileInfoIndex].GetFilename().mid(lastSlash + 1);
+	unsigned int renameIndex = 0;
+
+	// Check for name clashes
+	for (int i = 0; i < fileInfoIndex; i++)
+	{
+		lastSlash = fileInfos[i].GetFilename().lastIndexOf('/');
+
+		if (lastSlash < 0)
+			lastSlash = fileInfos[i].GetFilename().lastIndexOf('\\');
+
+		QString otherFilename = fileInfos[i].GetFilename().mid(lastSlash + 1);
+
+		if (filename == otherFilename)
+			renameIndex++;
+	}
+
+	if (renameIndex > 0)
+	{
+		int lastPeriodIndex = filename.lastIndexOf(QChar('.'));
+		QString shortFilename;
+		QString fileType;
+
+		if (lastPeriodIndex >= 0)
+		{
+			shortFilename = filename.left(lastPeriodIndex);
+			fileType = filename.mid(lastPeriodIndex);
+		}
+		else
+		{
+			shortFilename = filename;
+		}
+
+		unsigned int renameIndexOffset = 0;
+		bool validIndexOffset = true;
+
+		// Before we append a rename index we must ensure it doesn't produce further collisions.
+		for (int i = 0; i < fileInfos.length(); i++)
+		{
+			int lastSlash = fileInfos[i].GetFilename().lastIndexOf('/');
+
+			if (lastSlash < 0)
+				lastSlash = fileInfos[i].GetFilename().lastIndexOf('\\');
+
+			QString otherFilename = fileInfos[i].GetFilename().mid(lastSlash + 1);
+
+			if (otherFilename.length() > filename.length() + 1)
+			{
+				QString trimmedOtherFilename = otherFilename.left(shortFilename.length());
+
+				if (shortFilename == trimmedOtherFilename)
+				{
+					lastPeriodIndex = otherFilename.lastIndexOf(QChar('.'));
+					QString shortOtherFilename;
+
+					if (lastPeriodIndex >= 0)
+						shortOtherFilename = otherFilename.left(lastPeriodIndex);
+					else
+						shortOtherFilename = otherFilename;
+
+					QRegExp renameExp("-[0-9]+");
+					
+					if (renameExp.lastIndexIn(shortOtherFilename) == shortFilename.length())
+					{
+						unsigned int trailingInteger = shortOtherFilename.mid(shortFilename.length() + 1).toUInt(&validIndexOffset);
+
+						if (!validIndexOffset)
+							break;
+
+						if (trailingInteger > renameIndexOffset)
+							renameIndexOffset = trailingInteger;
+					}
+				}
+			}
+		}
+
+		if (validIndexOffset)
+		{
+			// Ensure renaming won't involve integer overflow!
+			if (renameIndex > static_cast<unsigned int>(-1) - renameIndexOffset)
+				validIndexOffset = false;
+		}
+
+		if (validIndexOffset)
+		{
+			shortFilename.append(QChar('-'));
+			shortFilename.append(QString::number(renameIndexOffset + renameIndex));
+
+			return (shortFilename + fileType);
+		}
+		else
+		{
+			// Fallback behaviour... an absolutely terrible brute force implementation!
+			QString filename;
+			QString renamePrefix;
+
+			for (;;)
+			{
+				renamePrefix.append(QChar('+'));
+
+				for (unsigned int i = 0; i < static_cast<unsigned int>(-1); i++)
+				{
+					filename = shortFilename + renamePrefix + QString::number(i) + fileType;
+
+					bool valid = true;
+
+					for (int i = 0; i < fileInfos.length(); i++)
+					{
+						int lastSlash = fileInfos[i].GetFilename().lastIndexOf('/');
+
+						if (lastSlash < 0)
+							lastSlash = fileInfos[i].GetFilename().lastIndexOf('\\');
+
+						if (filename == fileInfos[i].GetFilename().mid(lastSlash + 1))
+						{
+							valid = false;
+							break;
+						}
+					}
+
+					if (valid)
+						return (filename);
+				}
+			}
+		}
+	}
+	else
+	{
+		return (filename);
+	}
+}
+
+QString Packaging::ClashlessFilename(const QList<FileInfo>& fileInfos, const QString& filename)
+{
+	unsigned int renameIndex = 0;
+
+	// Check for name clashes
+	for (int i = 0; i < fileInfos.length(); i++)
+	{
+		int lastSlash = fileInfos[i].GetFilename().lastIndexOf('/');
+
+		if (lastSlash < 0)
+			lastSlash = fileInfos[i].GetFilename().lastIndexOf('\\');
+
+		QString otherFilename = fileInfos[i].GetFilename().mid(lastSlash + 1);
+
+		if (filename == otherFilename)
+			renameIndex++;
+	}
+
+	if (renameIndex > 0)
+	{
+		int lastPeriodIndex = filename.lastIndexOf(QChar('.'));
+		QString shortFilename;
+		QString fileType;
+
+		if (lastPeriodIndex >= 0)
+		{
+			shortFilename = filename.left(lastPeriodIndex);
+			fileType = filename.mid(lastPeriodIndex);
+		}
+		else
+		{
+			shortFilename = filename;
+		}
+
+		unsigned int renameIndexOffset = 0;
+		bool validIndexOffset = true;
+
+		// Before we append a rename index we must ensure it doesn't produce further collisions.
+		for (int i = 0; i < fileInfos.length(); i++)
+		{
+			int lastSlash = fileInfos[i].GetFilename().lastIndexOf('/');
+
+			if (lastSlash < 0)
+				lastSlash = fileInfos[i].GetFilename().lastIndexOf('\\');
+
+			QString otherFilename = fileInfos[i].GetFilename().mid(lastSlash + 1);
+
+			if (otherFilename.length() > filename.length() + 1)
+			{
+				QString trimmedOtherFilename = otherFilename.left(filename.length());
+
+				if (filename == trimmedOtherFilename)
+				{
+					lastPeriodIndex = otherFilename.lastIndexOf(QChar('.'));
+					QString shortOtherFilename;
+
+					if (lastPeriodIndex >= 0)
+						shortOtherFilename = otherFilename.left(lastPeriodIndex);
+					else
+						shortOtherFilename = otherFilename;
+
+					QRegExp renameExp("-[0-9]+");
+					
+					if (renameExp.lastIndexIn(shortOtherFilename) == shortFilename.length())
+					{
+						unsigned int trailingInteger = shortOtherFilename.mid(shortFilename.length() + 1).toUInt(&validIndexOffset);
+
+						if (!validIndexOffset)
+							break;
+
+						if (trailingInteger > renameIndexOffset)
+							renameIndexOffset = trailingInteger;
+					}
+				}
+			}
+		}
+
+		if (validIndexOffset)
+		{
+			// Ensure renaming won't involve integer overflow!
+			if (renameIndex > static_cast<unsigned int>(-1) - renameIndexOffset)
+				validIndexOffset = false;
+		}
+
+		if (validIndexOffset)
+		{
+			shortFilename.append(QChar('-'));
+			shortFilename.append(QString::number(renameIndexOffset + renameIndex));
+
+			return (shortFilename + fileType);
+		}
+		else
+		{
+			// Fallback behaviour, brute-force/semi-random.
+			bool valid;
+			QString filename;
+
+			do
+			{
+				valid = true;
+
+				filename = shortFilename + "-";
+				for (int i = 0; i < 8; i++)
+					filename.append(QChar(qrand() % ('Z' - 'A' + 1) + 'A'));
+
+				for (int i = 0; i < fileInfos.length(); i++)
+				{
+					int lastSlash = fileInfos[i].GetFilename().lastIndexOf('/');
+
+					if (lastSlash < 0)
+						lastSlash = fileInfos[i].GetFilename().lastIndexOf('\\');
+
+					if (filename == fileInfos[i].GetFilename().mid(lastSlash + 1))
+					{
+						valid = false;
+						break;
+					}
+				}
+			} while (!valid);
+
+			return (filename);
+		}
+	}
+	else
+	{
+		return (filename);
+	}
+}
diff --git a/heimdall-frontend/Source/Packaging.h b/heimdall-frontend/Source/Packaging.h
index 402f786..b6bac8d 100644
--- a/heimdall-frontend/Source/Packaging.h
+++ b/heimdall-frontend/Source/Packaging.h
@@ -103,7 +103,7 @@ namespace HeimdallFrontend
 			// TODO: Add support for sparse files to both methods?
 			static bool ExtractTar(QTemporaryFile& tarFile, PackageData *packageData);
 
-			static bool WriteTarEntry(const QString& filename, QTemporaryFile *tarFile, bool firmwareXml = false);
+			static bool WriteTarEntry(const QString& filePath, QTemporaryFile *tarFile, const QString& entryFilename);
 			static bool CreateTar(const FirmwareInfo& firmwareInfo, QTemporaryFile *tarFile); // Uses original TAR format.
 
 		public:
@@ -112,6 +112,9 @@ namespace HeimdallFrontend
 
 			static bool ExtractPackage(const QString& packagePath, PackageData *packageData);
 			static bool BuildPackage(const QString& packagePath, const FirmwareInfo& firmwareInfo);
+
+			static QString ClashlessFilename(const QList<FileInfo>& fileInfos, int fileInfoIndex);
+			static QString ClashlessFilename(const QList<FileInfo>& fileInfos, const QString& filename);
 	};
 }
 
diff --git a/heimdall-frontend/heimdall-frontend.vcxproj b/heimdall-frontend/heimdall-frontend.vcxproj
index 87be4e2..365e6e3 100644
--- a/heimdall-frontend/heimdall-frontend.vcxproj
+++ b/heimdall-frontend/heimdall-frontend.vcxproj
@@ -37,6 +37,9 @@
     <Import Project="$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props" Condition="exists('$(UserRootDir)\Microsoft.Cpp.$(Platform).user.props')" Label="LocalAppDataPlatform" />

     <Import Project="Qt4VSPropertySheet.props" />

   </ImportGroup>

+  <ImportGroup Label="PropertySheets" Condition="'$(Configuration)|$(Platform)'=='Release - Static MSVC|Win32'">

+    <Import Project="Qt4VSPropertySheet.props" />

+  </ImportGroup>

   <PropertyGroup Label="UserMacros" />

   <PropertyGroup>

     <_ProjectFileVersion>10.0.30319.1</_ProjectFileVersion>

@@ -100,7 +103,6 @@
   <ItemGroup>

     <ClCompile Include="GeneratedFiles\Debug\moc_mainwindow.cpp">

       <ExcludedFromBuild Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">true</ExcludedFromBuild>

-      <ExcludedFromBuild Condition="'$(Configuration)|$(Platform)'=='Release - Static MSVC|Win32'">true</ExcludedFromBuild>

     </ClCompile>

     <ClCompile Include="GeneratedFiles\qrc_mainwindow.cpp">

       <PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">

@@ -110,7 +112,6 @@
     </ClCompile>

     <ClCompile Include="GeneratedFiles\Release\moc_mainwindow.cpp">

       <ExcludedFromBuild Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">true</ExcludedFromBuild>

-      <ExcludedFromBuild Condition="'$(Configuration)|$(Platform)'=='Release - Static MSVC|Win32'">true</ExcludedFromBuild>

     </ClCompile>

     <ClCompile Include="Source\aboutform.cpp" />

     <ClCompile Include="Source\Alerts.cpp" />

diff --git a/heimdall-frontend/mainwindow.ui b/heimdall-frontend/mainwindow.ui
index 985fe10..984f4ba 100644
--- a/heimdall-frontend/mainwindow.ui
+++ b/heimdall-frontend/mainwindow.ui
@@ -130,9 +130,9 @@
      <widget class="QGroupBox" name="platformGroup">

       <property name="geometry">

        <rect>

-        <x>350</x>

+        <x>340</x>

         <y>80</y>

-        <width>151</width>

+        <width>161</width>

         <height>61</height>

        </rect>

       </property>

@@ -147,7 +147,7 @@
         <rect>

          <x>10</x>

          <y>30</y>

-         <width>131</width>

+         <width>141</width>

          <height>21</height>

         </rect>

        </property>

@@ -216,7 +216,7 @@
        <rect>

         <x>10</x>

         <y>80</y>

-        <width>221</width>

+        <width>211</width>

         <height>61</height>

        </rect>

       </property>

@@ -231,7 +231,7 @@
         <rect>

          <x>10</x>

          <y>30</y>

-         <width>201</width>

+         <width>191</width>

          <height>21</height>

         </rect>

        </property>

@@ -288,7 +288,7 @@
      <widget class="QGroupBox" name="versionGroup">

       <property name="geometry">

        <rect>

-        <x>240</x>

+        <x>230</x>

         <y>80</y>

         <width>101</width>

         <height>61</height>

@@ -643,9 +643,9 @@
         </property>

         <property name="geometry">

          <rect>

-          <x>120</x>

+          <x>140</x>

           <y>30</y>

-          <width>261</width>

+          <width>241</width>

           <height>22</height>

          </rect>

         </property>

@@ -655,7 +655,7 @@
          <rect>

           <x>10</x>

           <y>30</y>

-          <width>101</width>

+          <width>121</width>

           <height>16</height>

          </rect>

         </property>

@@ -668,7 +668,7 @@
          <rect>

           <x>10</x>

           <y>60</y>

-          <width>101</width>

+          <width>121</width>

           <height>16</height>

          </rect>

         </property>

@@ -727,9 +727,9 @@
         </property>

         <property name="geometry">

          <rect>

-          <x>120</x>

+          <x>140</x>

           <y>60</y>

-          <width>261</width>

+          <width>241</width>

           <height>21</height>

          </rect>

         </property>

@@ -885,7 +885,7 @@
        <rect>

         <x>0</x>

         <y>230</y>

-        <width>491</width>

+        <width>471</width>

         <height>231</height>

        </rect>

       </property>

@@ -897,7 +897,7 @@
         <rect>

          <x>10</x>

          <y>30</y>

-         <width>471</width>

+         <width>451</width>

          <height>161</height>

         </rect>

        </property>

@@ -908,9 +908,9 @@
        </property>

        <property name="geometry">

         <rect>

-         <x>350</x>

+         <x>320</x>

          <y>200</y>

-         <width>131</width>

+         <width>141</width>

          <height>23</height>

         </rect>

        </property>

@@ -1024,9 +1024,9 @@
        </property>

        <property name="geometry">

         <rect>

-         <x>550</x>

+         <x>560</x>

          <y>20</y>

-         <width>191</width>

+         <width>181</width>

          <height>81</height>

         </rect>

        </property>

@@ -1034,7 +1034,7 @@
       <widget class="QGroupBox" name="createDeveloperInfoGroup">

        <property name="geometry">

         <rect>

-         <x>300</x>

+         <x>310</x>

          <y>20</y>

          <width>241</width>

          <height>101</height>

@@ -1048,7 +1048,7 @@
          <rect>

           <x>10</x>

           <y>30</y>

-          <width>41</width>

+          <width>51</width>

           <height>16</height>

          </rect>

         </property>

@@ -1062,9 +1062,9 @@
         </property>

         <property name="geometry">

          <rect>

-          <x>60</x>

+          <x>70</x>

           <y>30</y>

-          <width>171</width>

+          <width>161</width>

           <height>21</height>

          </rect>

         </property>

@@ -1110,7 +1110,7 @@
         <rect>

          <x>10</x>

          <y>20</y>

-         <width>281</width>

+         <width>291</width>

          <height>101</height>

         </rect>

        </property>

@@ -1122,7 +1122,7 @@
          <rect>

           <x>10</x>

           <y>60</y>

-          <width>71</width>

+          <width>81</width>

           <height>16</height>

          </rect>

         </property>

@@ -1135,7 +1135,7 @@
          <rect>

           <x>10</x>

           <y>30</y>

-          <width>71</width>

+          <width>81</width>

           <height>16</height>

          </rect>

         </property>

@@ -1149,7 +1149,7 @@
         </property>

         <property name="geometry">

          <rect>

-          <x>90</x>

+          <x>100</x>

           <y>30</y>

           <width>181</width>

           <height>21</height>

@@ -1165,7 +1165,7 @@
         </property>

         <property name="geometry">

          <rect>

-          <x>90</x>

+          <x>100</x>

           <y>60</y>

           <width>181</width>

           <height>21</height>

@@ -1209,9 +1209,9 @@
      <widget class="QGroupBox" name="createDeviceInfoGroup">

       <property name="geometry">

        <rect>

-        <x>500</x>

+        <x>480</x>

         <y>240</y>

-        <width>271</width>

+        <width>291</width>

         <height>151</height>

        </rect>

       </property>

@@ -1223,7 +1223,7 @@
         <rect>

          <x>10</x>

          <y>60</y>

-         <width>91</width>

+         <width>111</width>

          <height>16</height>

         </rect>

        </property>

@@ -1237,7 +1237,7 @@
        </property>

        <property name="geometry">

         <rect>

-         <x>110</x>

+         <x>130</x>

          <y>60</y>

          <width>151</width>

          <height>21</height>

@@ -1252,7 +1252,7 @@
         <rect>

          <x>10</x>

          <y>30</y>

-         <width>91</width>

+         <width>111</width>

          <height>16</height>

         </rect>

        </property>

@@ -1266,7 +1266,7 @@
        </property>

        <property name="geometry">

         <rect>

-         <x>110</x>

+         <x>130</x>

          <y>30</y>

          <width>151</width>

          <height>21</height>

@@ -1281,7 +1281,7 @@
         <rect>

          <x>10</x>

          <y>90</y>

-         <width>91</width>

+         <width>111</width>

          <height>16</height>

         </rect>

        </property>

@@ -1295,7 +1295,7 @@
        </property>

        <property name="geometry">

         <rect>

-         <x>110</x>

+         <x>130</x>

          <y>90</y>

          <width>151</width>

          <height>21</height>

@@ -1311,9 +1311,9 @@
        </property>

        <property name="geometry">

         <rect>

-         <x>150</x>

+         <x>160</x>

          <y>120</y>

-         <width>111</width>

+         <width>121</width>

          <height>23</height>

         </rect>

        </property>

diff --git a/heimdall/source/main.cpp b/heimdall/source/main.cpp
index 63056bd..f5bf173 100644
--- a/heimdall/source/main.cpp
+++ b/heimdall/source/main.cpp
@@ -244,6 +244,35 @@ void closeFiles(map<string, FILE *> argumentfileMap)
 	argumentfileMap.clear();
 }
 
+bool getDeviceInfo(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);
+	}
+}
+
 int downloadPitFile(BridgeManager *bridgeManager, unsigned char **pitBuffer)
 {
 	Interface::Print("Downloading device's PIT file...\n");
@@ -345,32 +374,6 @@ bool flashFile(BridgeManager *bridgeManager, unsigned int partitionIndex, const
 bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFileMap, bool repartition)
 {
 	bool success;
-	
-	// ---------- 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);
-	}
 
 	// ------------- SEND TOTAL BYTES TO BE TRANSFERRED -------------
 
@@ -397,7 +400,7 @@ bool attemptFlash(BridgeManager *bridgeManager, map<string, FILE *> argumentFile
 
 	DeviceInfoResponse *deviceInfoResponse = new DeviceInfoResponse();
 	success = bridgeManager->ReceivePacket(deviceInfoResponse);
-	deviceInfoResult = deviceInfoResponse->GetUnknown();
+	int deviceInfoResult = deviceInfoResponse->GetUnknown();
 	delete deviceInfoResponse;
 
 	if (!success)
@@ -663,7 +666,7 @@ int main(int argc, char **argv)
 				return (0);
 			}
 
-			if (!bridgeManager->BeginSession())
+			if (!bridgeManager->BeginSession() || !getDeviceInfo(bridgeManager))
 			{
 				closeFiles(argumentFileMap);
 				delete bridgeManager;
@@ -683,7 +686,7 @@ int main(int argc, char **argv)
 
 		case Interface::kActionClosePcScreen:
 		{
-			if (!bridgeManager->BeginSession())
+			if (!bridgeManager->BeginSession() || !getDeviceInfo(bridgeManager))
 			{
 				delete bridgeManager;
 				return (-1);
@@ -710,7 +713,7 @@ int main(int argc, char **argv)
 				return (0);
 			}
 
-			if (!bridgeManager->BeginSession())
+			if (!bridgeManager->BeginSession() || !getDeviceInfo(bridgeManager))
 			{
 				delete bridgeManager;
 				fclose(outputPitFile);
@@ -761,7 +764,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())
+			if (!bridgeManager->BeginSession() || !getDeviceInfo(bridgeManager))
 			{
 				fclose(dumpFile);
 
@@ -780,7 +783,7 @@ int main(int argc, char **argv)
 
 		case Interface::kActionPrintPit:
 		{
-			if (!bridgeManager->BeginSession())
+			if (!bridgeManager->BeginSession() || !getDeviceInfo(bridgeManager))
 			{
 				delete bridgeManager;
 				return (-1);