Merge pull request #913 from avast/mpress-resources

unpacker/mpress: Properly copy non-packer related sections to the unpacked file
avast · Jan 8, 2021 · f4da380 · f4da380
2 parents da1c736 + aa15ad9
commit f4da380
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 8 deletions.
diff --git a/src/unpackertool/plugins/mpress/mpress.cpp b/src/unpackertool/plugins/mpress/mpress.cpp
@@ -658,17 +658,18 @@ void MpressPlugin::saveFile(const std::string& fileName, DynamicBuffer& content)
 	// Headers
 	imageLoader.Save(fileName.c_str());
 
+	std::fstream outputFile(fileName, std::ios::binary | std::ios::out | std::ios::in);
 	// Copy the section bytes from original file for the sections preceding the packed section
 	for (std::uint32_t index = 0; index < _packedContentSect->getSecSeg()->getIndex(); ++index)
-		copySectionFromOriginalFile(index, fileName, index);
+		copySectionFromOriginalFile(index, outputFile, index);
 
 	// Copy the section bytes in between packed section and EP section
 	for (std::uint32_t index = _packedContentSect->getSecSeg()->getIndex() + _addedSectionCount; index < _file->getEpSegment()->getSecSeg()->getIndex(); ++index)
-		copySectionFromOriginalFile(index, fileName, index + _addedSectionCount);
+		copySectionFromOriginalFile(index, outputFile, index + _addedSectionCount);
 
 	// Copy the section bytes from original file for sections after EP section excluded
 	for (std::uint32_t index = _file->getEpSegment()->getSecSeg()->getIndex() + 1; index < _file->getNumberOfSegments(); ++index)
-		copySectionFromOriginalFile(index, fileName, index + _addedSectionCount);
+		copySectionFromOriginalFile(index, outputFile, index + _addedSectionCount);
 
 	// Write content of new import section
 	std::uint32_t Rva = imageLoader.getDataDirRva(PeLib::PELIB_IMAGE_DIRECTORY_ENTRY_IMPORT);
@@ -691,18 +692,22 @@ void MpressPlugin::saveFile(const std::string& fileName, DynamicBuffer& content)
 
 	// Write the unpacked content to the packed content section
 	// Use regular file as we will write more sections at once
-	std::fstream outputFile(fileName, std::ios::binary | std::ios::out | std::ios::in);
 	outputFile.seekp(imageLoader.getSectionHeader(_packedContentSect->getSecSeg()->getIndex())->PointerToRawData, std::ios_base::beg);
 	outputFile.write(reinterpret_cast<const char*>(content.getRawBuffer()), content.getRealDataSize());
 	outputFile.close();
 }
 
-void MpressPlugin::copySectionFromOriginalFile(std::uint32_t origSectIndex, const std::string& newFileName, std::uint32_t newSectIndex)
+void MpressPlugin::copySectionFromOriginalFile(std::uint32_t origSectIndex, std::ostream& outputFile, std::uint32_t newSectIndex)
 {
 	const retdec::loader::Segment* seg = _file->getSegment(origSectIndex);
 	std::vector<std::uint8_t> bytes;
 	seg->getBytes(bytes);
-	//_peFile->peHeader().writeSectionData(newFileName, newSectIndex, bytes);
+
+	PeLib::ImageLoader & imageLoader = _peFile->imageLoader();
+	const auto* newSect = imageLoader.getSectionHeader(newSectIndex);
+	outputFile.seekp(newSect->PointerToRawData, std::ios_base::beg);
+	outputFile.write(reinterpret_cast<const char*>(bytes.data()), std::min(static_cast<std::uint32_t>(bytes.size()), newSect->SizeOfRawData));
+
 }
 
 } // namespace mpress

diff --git a/src/unpackertool/plugins/mpress/mpress.h b/src/unpackertool/plugins/mpress/mpress.h
@@ -86,7 +86,7 @@ class MpressPlugin : public Plugin
 	MpressUnpackerStub detectUnpackerStubVersion();
 	MpressFixStub detectFixStubVersion(retdec::utils::DynamicBuffer& unpackedContent);
 	void saveFile(const std::string& fileName, retdec::utils::DynamicBuffer& content);
-	void copySectionFromOriginalFile(std::uint32_t origSectIndex, const std::string& newFileName, std::uint32_t newSectIndex);
+	void copySectionFromOriginalFile(std::uint32_t origSectIndex, std::ostream& outputFile, std::uint32_t newSectIndex);
 
 	std::unique_ptr<retdec::loader::Image> _file;
 	PeLib::PeFileT * _peFile;

diff --git a/src/unpackertool/plugins/upx/pe/pe_upx_stub.cpp b/src/unpackertool/plugins/upx/pe/pe_upx_stub.cpp
@@ -1238,7 +1238,7 @@ template <int bits> void PeUpxStub<bits>::saveFile(const std::string& outputFile
 	if (!_coffSymbolTable.empty())
 		retdec::utils::writeFile(outputFileHandle, _coffSymbolTable, imageLoader.getPointerToSymbolTable());
 	outputFileHandle.close();
-	
+
 	// Write resources at the end, because they would be rewritten by unpackedData which have them zeroed
 	if((Rva = imageLoader.getDataDirRva(PeLib::PELIB_IMAGE_DIRECTORY_ENTRY_RESOURCE)) != 0)
 		_newPeFile->resDir().write(outputFile, imageLoader.getFileOffsetFromRva(Rva), Rva);