Skip to content

Commit

Permalink
Previous behavior of ImageLoader::Save() is back as special case (nee…
Browse files Browse the repository at this point in the history 
…ded for unpackers)
  • Loading branch information
Ladislav Zezula committed Oct 9, 2021
1 parent a00a663 commit 6753eae
Show file tree
Hide file tree
Showing 5 changed files with 234 additions and 37 deletions.
24 changes: 12 additions & 12 deletions include/retdec/pelib/ImageLoader.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,16 +20,6 @@

namespace PeLib {

//-----------------------------------------------------------------------------
// Enum for ImageLoader::Load()

enum FileFlags : std::uint32_t
{
Default = 0,
HeadersOnly = 1,
LoadAsImage = 2
};

//-----------------------------------------------------------------------------
// Enum for ImageLoader::getFieldOffset()

Expand Down Expand Up @@ -71,6 +61,13 @@ const std::uint32_t BuildNumber10 = 10240; // Behavior equal to Windows
const std::uint32_t BuildNumberMask = 0x0FFFF; // Mask for extracting the operating system
const std::uint32_t BuildNumber64Bit = 0x10000; // Emulate 64-bit system

//-----------------------------------------------------------------------------
// Flags for ImageLoader::Load() and ImageLoader::Save()

const std::uint32_t IoFlagHeadersOnly = 1; // Only load/save PE headers
const std::uint32_t IoFlagNewFile = 2; // Create the PE as new file (for unpackers)
const std::uint32_t IoFlagLoadAsImage = 4; // Load the data as mapped image file

//-----------------------------------------------------------------------------
// Structure for comparison with Windows mapped images

Expand Down Expand Up @@ -154,8 +151,8 @@ class ImageLoader
int Load(std::istream & fs, std::streamoff fileOffset = 0, bool loadHeadersOnly = false);
int Load(const char * fileName, bool loadHeadersOnly = false);

int Save(std::ostream & fs, std::streamoff fileOffset = 0, FileFlags saveFlags = FileFlags::Default);
int Save(const char * fileName, FileFlags saveFlags = FileFlags::Default);
int Save(std::ostream & fs, std::streamoff fileOffset = 0, std::uint32_t saveFlags = 0);
int Save(const char * fileName, std::uint32_t saveFlags = 0);

bool relocateImage(std::uint64_t newImageBase);

Expand Down Expand Up @@ -404,11 +401,14 @@ class ImageLoader

int captureDosHeader(ByteBuffer & fileData);
int saveToFile(std::ostream & fs, std::streamoff fileOffset, std::size_t rva, std::size_t length);
int saveDosHeaderNew(std::ostream & fs, std::streamoff fileOffset);
int saveDosHeader(std::ostream & fs, std::streamoff fileOffset);
int captureNtHeaders(ByteBuffer & fileData);
int saveNtHeadersNew(std::ostream & fs, std::streamoff fileOffset);
int saveNtHeaders(std::ostream & fs, std::streamoff fileOffset);
int captureSectionName(ByteBuffer & fileData, std::string & sectionName, const std::uint8_t * name);
int captureSectionHeaders(ByteBuffer & fileData);
int saveSectionHeadersNew(std::ostream & fs, std::streamoff fileOffset);
int saveSectionHeaders(std::ostream & fs, std::streamoff fileOffset);
int captureImageSections(ByteBuffer & fileData);
int captureOptionalHeader32(std::uint8_t * fileData, std::uint8_t * filePtr, std::uint8_t * fileEnd);
Expand Down
1 change: 1 addition & 0 deletions src/bin2llvmir/optimizations/decoder/decoder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,7 @@ void Decoder::decodeJumpTarget(const JumpTarget& jt)
LOG << "\t\t\t" << "translating = " << addr << std::endl;

Address oldAddr = addr;

auto res = translate(bytes, addr, irb);

if (res.failed() || res.llvmInsn == nullptr)
Expand Down
242 changes: 219 additions & 23 deletions src/pelib/ImageLoader.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -912,7 +912,6 @@ PeLib::LoaderError PeLib::ImageLoader::loaderError() const

//-----------------------------------------------------------------------------
// Interface for loading files

int PeLib::ImageLoader::Load(
ByteBuffer & fileData,
bool loadHeadersOnly)
Expand Down Expand Up @@ -1053,35 +1052,85 @@ int PeLib::ImageLoader::Load(
int PeLib::ImageLoader::Save(
std::ostream & fs,
std::streamoff fileOffset,
FileFlags saveFlags)
std::uint32_t saveFlags)
{
int fileError;

// Check and capture DOS header
fileError = saveDosHeader(fs, fileOffset);
if(fileError != ERROR_NONE)
return fileError;
// This save mode is intended for unpackers. Headers are constructed
// from metadata and sections are filled with zeros
if(saveFlags & IoFlagNewFile)
{
// Save the DOS header
fileError = saveDosHeaderNew(fs, fileOffset);
if(fileError != ERROR_NONE)
return fileError;

// Check and capture NT headers
fileError = saveNtHeaders(fs, fileOffset + dosHeader.e_lfanew);
if(fileError != ERROR_NONE)
return fileError;
// Save the NT headers
fileError = saveNtHeadersNew(fs, fileOffset + dosHeader.e_lfanew);
if(fileError != ERROR_NONE)
return fileError;

// Check and capture section headers
fileOffset = fileOffset + dosHeader.e_lfanew + sizeof(PELIB_IMAGE_NT_SIGNATURE) + sizeof(PELIB_IMAGE_FILE_HEADER) + fileHeader.SizeOfOptionalHeader;
fileError = saveSectionHeaders(fs, fileOffset);
if(fileError != ERROR_NONE)
return fileError;
// Check and capture section headers
fileOffset = fileOffset + dosHeader.e_lfanew + sizeof(PELIB_IMAGE_NT_SIGNATURE) + sizeof(PELIB_IMAGE_FILE_HEADER) + fileHeader.SizeOfOptionalHeader;
fileError = saveSectionHeadersNew(fs, fileOffset);
if(fileError != ERROR_NONE)
return fileError;

// Write section data to the file, up to size of image
if(!(saveFlags & IoFlagHeadersOnly))
{
// Get the curent file offset and file size
fileOffset += sections.size() * sizeof(PELIB_IMAGE_SECTION_HEADER);
std::streamoff fileSize = fileOffset;

// Estimate the file size with data
for(const auto & section : sections)
{
if(section.SizeOfRawData != 0)
{
if((section.PointerToRawData + section.SizeOfRawData) > fileSize)
fileSize = section.PointerToRawData + section.SizeOfRawData;
}
}

// Shall we write data to the file?
if(fileSize > fileOffset)
{
std::vector<char> ZeroBuffer(fileSize - fileOffset);

// Write zeros to the rest of the file, up to size of image
if(!(saveFlags & FileFlags::HeadersOnly))
fs.seekp(fileOffset, std::ios::beg);
fs.write(ZeroBuffer.data(), ZeroBuffer.size());
}
}
}
else
{
// Write each section
for(const auto & section : sections)
// Save the DOS header
fileError = saveDosHeader(fs, fileOffset);
if(fileError != ERROR_NONE)
return fileError;

// Save the NT headers
fileError = saveNtHeaders(fs, fileOffset + dosHeader.e_lfanew);
if(fileError != ERROR_NONE)
return fileError;

// Check and capture section headers
fileOffset = fileOffset + dosHeader.e_lfanew + sizeof(PELIB_IMAGE_NT_SIGNATURE) + sizeof(PELIB_IMAGE_FILE_HEADER) + fileHeader.SizeOfOptionalHeader;
fileError = saveSectionHeaders(fs, fileOffset);
if(fileError != ERROR_NONE)
return fileError;

// Write section data to the file, up to size of image
if(!(saveFlags & IoFlagHeadersOnly))
{
fileError = saveToFile(fs, section.PointerToRawData, section.VirtualAddress, section.SizeOfRawData);
if(fileError != ERROR_NONE)
return fileError;
// Write each section
for(const auto & section : sections)
{
fileError = saveToFile(fs, section.PointerToRawData, section.VirtualAddress, section.SizeOfRawData);
if(fileError != ERROR_NONE)
return fileError;
}
}
}

Expand All @@ -1090,7 +1139,7 @@ int PeLib::ImageLoader::Save(

int PeLib::ImageLoader::Save(
const char * fileName,
FileFlags saveFlags)
std::uint32_t saveFlags)
{
std::ofstream fs(fileName, std::ifstream::out | std::ifstream::binary);
if(!fs.is_open())
Expand Down Expand Up @@ -1611,6 +1660,16 @@ int PeLib::ImageLoader::saveToFile(
return ERROR_NONE;
}

int PeLib::ImageLoader::saveDosHeaderNew(
std::ostream & fs,
std::streamoff fileOffset)
{
// Write DOS header as-is
fs.seekp(fileOffset, std::ios::beg);
fs.write(reinterpret_cast<char *>(&dosHeader), sizeof(PELIB_IMAGE_DOS_HEADER));
return ERROR_NONE;
}

int PeLib::ImageLoader::saveDosHeader(
std::ostream & fs,
std::streamoff fileOffset)
Expand Down Expand Up @@ -1753,6 +1812,117 @@ int PeLib::ImageLoader::captureNtHeaders(ByteBuffer & fileData)
return ERROR_NONE;
}

int PeLib::ImageLoader::saveNtHeadersNew(
std::ostream & fs,
std::streamoff fileOffset)
{
// Calculate the size of the optional header. Any version of PE file,
// 32 or 64-bit, must have this field set to a correct value.
std::uint32_t sizeOfOptionalHeader = getFieldOffset(PELIB_MEMBER_TYPE::OPTHDR_sizeof_fixed) + optionalHeader.NumberOfRvaAndSizes * sizeof(PELIB_IMAGE_DATA_DIRECTORY);

// Move to the required file offset
fs.seekp(fileOffset, std::ios::beg);

// Write the NT signature
fs.write(reinterpret_cast<char *>(&ntSignature), sizeof(ntSignature));

// Write the file header
fileHeader.SizeOfOptionalHeader = sizeOfOptionalHeader;
fileHeader.NumberOfSections = (std::uint16_t)sections.size();
fs.write(reinterpret_cast<char *>(&fileHeader), sizeof(PELIB_IMAGE_FILE_HEADER));

// Write the optional header. Note that we need to distinguish 32-bit and 64-bit header
if(optionalHeader.Magic == PELIB_IMAGE_NT_OPTIONAL_HDR32_MAGIC)
{
PELIB_IMAGE_OPTIONAL_HEADER32 optionalHeader32;

// Verify some of the data to make sure they are able to convert to 32-bit values
if((optionalHeader.ImageBase >> 0x20) != 0)
return ERROR_INVALID_FILE;
if((optionalHeader.SizeOfStackReserve >> 0x20) != 0)
return ERROR_INVALID_FILE;
if((optionalHeader.SizeOfHeapReserve >> 0x20) != 0)
return ERROR_INVALID_FILE;

// Convert the optional header to 32-bit variant
optionalHeader32.Magic = optionalHeader.Magic;
optionalHeader32.MajorLinkerVersion = optionalHeader.MajorLinkerVersion;
optionalHeader32.MinorLinkerVersion = optionalHeader.MinorLinkerVersion;
optionalHeader32.SizeOfCode = optionalHeader.SizeOfCode;
optionalHeader32.SizeOfInitializedData = optionalHeader.SizeOfInitializedData;
optionalHeader32.SizeOfUninitializedData = optionalHeader.SizeOfUninitializedData;
optionalHeader32.AddressOfEntryPoint = optionalHeader.AddressOfEntryPoint;
optionalHeader32.BaseOfCode = optionalHeader.BaseOfCode;
optionalHeader32.BaseOfData = optionalHeader.BaseOfData;
optionalHeader32.ImageBase = (std::uint32_t)optionalHeader.ImageBase;
optionalHeader32.SectionAlignment = optionalHeader.SectionAlignment;
optionalHeader32.FileAlignment = optionalHeader.FileAlignment;
optionalHeader32.MajorOperatingSystemVersion = optionalHeader.MajorOperatingSystemVersion;
optionalHeader32.MinorOperatingSystemVersion = optionalHeader.MinorOperatingSystemVersion;
optionalHeader32.MajorImageVersion = optionalHeader.MajorImageVersion;
optionalHeader32.MinorImageVersion = optionalHeader.MinorImageVersion;
optionalHeader32.MajorSubsystemVersion = optionalHeader.MajorSubsystemVersion;
optionalHeader32.MinorSubsystemVersion = optionalHeader.MinorSubsystemVersion;
optionalHeader32.Win32VersionValue = optionalHeader.Win32VersionValue;
optionalHeader32.SizeOfImage = optionalHeader.SizeOfImage;
optionalHeader32.SizeOfHeaders = optionalHeader.SizeOfHeaders;
optionalHeader32.CheckSum = optionalHeader.CheckSum;
optionalHeader32.Subsystem = optionalHeader.Subsystem;
optionalHeader32.DllCharacteristics = optionalHeader.DllCharacteristics;
optionalHeader32.SizeOfStackReserve = (std::uint32_t)optionalHeader.SizeOfStackReserve;
optionalHeader32.SizeOfStackCommit = (std::uint32_t)optionalHeader.SizeOfStackCommit;
optionalHeader32.SizeOfHeapReserve = (std::uint32_t)optionalHeader.SizeOfHeapReserve;
optionalHeader32.SizeOfHeapCommit = (std::uint32_t)optionalHeader.SizeOfHeapCommit;
optionalHeader32.LoaderFlags = optionalHeader.LoaderFlags;
optionalHeader32.NumberOfRvaAndSizes = optionalHeader.NumberOfRvaAndSizes;
memcpy(&optionalHeader32.DataDirectory, &optionalHeader.DataDirectory, sizeof(optionalHeader.DataDirectory));

// Write to file
fs.write(reinterpret_cast<char *>(&optionalHeader32), sizeOfOptionalHeader);
}
else
{
PELIB_IMAGE_OPTIONAL_HEADER64 optionalHeader64;

// Convert the optional header to 64-bit variant
optionalHeader64.Magic = optionalHeader.Magic;
optionalHeader64.MajorLinkerVersion = optionalHeader.MajorLinkerVersion;
optionalHeader64.MinorLinkerVersion = optionalHeader.MinorLinkerVersion;
optionalHeader64.SizeOfCode = optionalHeader.SizeOfCode;
optionalHeader64.SizeOfInitializedData = optionalHeader.SizeOfInitializedData;
optionalHeader64.SizeOfUninitializedData = optionalHeader.SizeOfUninitializedData;
optionalHeader64.AddressOfEntryPoint = optionalHeader.AddressOfEntryPoint;
optionalHeader64.BaseOfCode = optionalHeader.BaseOfCode;
optionalHeader64.ImageBase = optionalHeader.ImageBase;
optionalHeader64.SectionAlignment = optionalHeader.SectionAlignment;
optionalHeader64.FileAlignment = optionalHeader.FileAlignment;
optionalHeader64.MajorOperatingSystemVersion = optionalHeader.MajorOperatingSystemVersion;
optionalHeader64.MinorOperatingSystemVersion = optionalHeader.MinorOperatingSystemVersion;
optionalHeader64.MajorImageVersion = optionalHeader.MajorImageVersion;
optionalHeader64.MinorImageVersion = optionalHeader.MinorImageVersion;
optionalHeader64.MajorSubsystemVersion = optionalHeader.MajorSubsystemVersion;
optionalHeader64.MinorSubsystemVersion = optionalHeader.MinorSubsystemVersion;
optionalHeader64.Win32VersionValue = optionalHeader.Win32VersionValue;
optionalHeader64.SizeOfImage = optionalHeader.SizeOfImage;
optionalHeader64.SizeOfHeaders = optionalHeader.SizeOfHeaders;
optionalHeader64.CheckSum = optionalHeader.CheckSum;
optionalHeader64.Subsystem = optionalHeader.Subsystem;
optionalHeader64.DllCharacteristics = optionalHeader.DllCharacteristics;
optionalHeader64.SizeOfStackReserve = optionalHeader.SizeOfStackReserve;
optionalHeader64.SizeOfStackCommit = optionalHeader.SizeOfStackCommit;
optionalHeader64.SizeOfHeapReserve = optionalHeader.SizeOfHeapReserve;
optionalHeader64.SizeOfHeapCommit = optionalHeader.SizeOfHeapCommit;
optionalHeader64.LoaderFlags = optionalHeader.LoaderFlags;
optionalHeader64.NumberOfRvaAndSizes = optionalHeader.NumberOfRvaAndSizes;
memcpy(&optionalHeader64.DataDirectory, &optionalHeader.DataDirectory, sizeof(optionalHeader64.DataDirectory));

// Write to file
fs.write(reinterpret_cast<char *>(&optionalHeader64), sizeOfOptionalHeader);
}

return ERROR_NONE;
}

int PeLib::ImageLoader::saveNtHeaders(
std::ostream & fs,
std::streamoff fileOffset)
Expand Down Expand Up @@ -1984,6 +2154,32 @@ int PeLib::ImageLoader::captureSectionHeaders(ByteBuffer & fileData)
return ERROR_NONE;
}

int PeLib::ImageLoader::saveSectionHeadersNew(
std::ostream & fs,
std::streamoff fileOffset)
{
PELIB_IMAGE_SECTION_HEADER * pHeaders;
std::size_t sectionCount = sections.size();
std::size_t index = 0;

if((pHeaders = new PELIB_IMAGE_SECTION_HEADER[sectionCount]) != nullptr)
{
// Populate the array with section headers
for(const auto & section : sections)
{
memcpy(pHeaders + index, section.Name, sizeof(PELIB_IMAGE_SECTION_HEADER));
index++;
}

// Write the section headers to file
fs.seekp(fileOffset, std::ios::beg);
fs.write(reinterpret_cast<char *>(pHeaders), sectionCount * sizeof(PELIB_IMAGE_SECTION_HEADER));
delete[] pHeaders;
}

return ERROR_NONE;
}

int PeLib::ImageLoader::saveSectionHeaders(
std::ostream & fs,
std::streamoff fileOffset)
Expand Down
2 changes: 1 addition & 1 deletion src/unpackertool/plugins/mpress/mpress.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -656,7 +656,7 @@ void MpressPlugin::saveFile(const std::string& fileName, DynamicBuffer& content)
std::remove(fileName.c_str());

// Headers
imageLoader.Save(fileName.c_str());
imageLoader.Save(fileName.c_str(), PeLib::IoFlagNewFile);

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
Expand Down
2 changes: 1 addition & 1 deletion src/unpackertool/plugins/upx/pe/pe_upx_stub.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1206,7 +1206,7 @@ template <int bits> void PeUpxStub<bits>::saveFile(const std::string& outputFile

// Write the DOS header, PE headers and section headers
pSectionHeader = imageLoader.getSectionHeader(_upx0Sect->getSecSeg()->getIndex());
imageLoader.Save(outputFile.c_str());
imageLoader.Save(outputFile.c_str(), PeLib::IoFlagNewFile);

// Save the import directory
if((Rva = imageLoader.getDataDirRva(PeLib::PELIB_IMAGE_DIRECTORY_ENTRY_IMPORT)) != 0)
Expand Down

0 comments on commit 6753eae

Please sign in to comment.