ImportRebuilder.cpp
No OneTemporary
Actions

Size

8 KB

Subscribers

None

ImportRebuilder.cpp
View Options


	#include "ImportRebuilder.h"
	#include "Scylla.h"
	#include "StringConversion.h"

	//#define DEBUG_COMMENTS


	bool ImportRebuilder::rebuildImportTable(const WCHAR * newFilePath, std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	bool retValue = false;

	if (isValidPeFile())
	{
	if (readPeSectionsFromFile())
	{
	setDefaultFileAlignment();

	retValue = buildNewImportTable(moduleList);

	if (retValue)
	{
	alignAllSectionHeaders();
	fixPeHeader();
	retValue = savePeFileToDisk(newFilePath);
	}
	}
	}

	return retValue;
	}

	bool ImportRebuilder::buildNewImportTable(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	createNewImportSection(moduleList);

	importSectionIndex = listPeSection.size() - 1;

	DWORD dwSize = fillImportSection(moduleList);

	if (!dwSize)
	{
	return false;
	}

	setFlagToIATSection((*moduleList.begin()).second.firstThunk);

	if (isPE32())
	{
	pNTHeader32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress = listPeSection[importSectionIndex].sectionHeader.VirtualAddress;
	pNTHeader32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size = (DWORD)(numberOfImportDescriptors * sizeof(IMAGE_IMPORT_DESCRIPTOR));
	}
	else
	{
	pNTHeader64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress = listPeSection[importSectionIndex].sectionHeader.VirtualAddress;
	pNTHeader64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size = (DWORD)(numberOfImportDescriptors * sizeof(IMAGE_IMPORT_DESCRIPTOR));
	}


	return true;
	}

	bool ImportRebuilder::createNewImportSection(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	char sectionName[IMAGE_SIZEOF_SHORT_NAME + 1] = {0};

	const WCHAR * sectionNameW = Scylla::config[IAT_SECTION_NAME].getString();

	calculateImportSizes(moduleList);

	if (wcslen(sectionNameW) > IMAGE_SIZEOF_SHORT_NAME)
	{
	strcpy_s(sectionName, ".SCY");
	}
	else
	{
	StringConversion::ToASCII(sectionNameW, sectionName, _countof(sectionName));
	}

	return addNewLastSection(sectionName, (DWORD)sizeOfImportSection, 0);
	}

	void ImportRebuilder::setFlagToIATSection(DWORD_PTR iatAddress)
	{
	for (size_t i = 0; i < listPeSection.size(); i++)
	{
	if ((listPeSection[i].sectionHeader.VirtualAddress <= iatAddress) && ((listPeSection[i].sectionHeader.VirtualAddress + listPeSection[i].sectionHeader.Misc.VirtualSize) > iatAddress))
	{
	//section must be read and writeable
	listPeSection[i].sectionHeader.Characteristics \|= IMAGE_SCN_MEM_READ\|IMAGE_SCN_MEM_WRITE;
	}
	}
	}

	DWORD ImportRebuilder::fillImportSection(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	std::map<DWORD_PTR, ImportModuleThunk>::iterator mapIt;
	std::map<DWORD_PTR, ImportThunk>::iterator mapIt2;
	PIMAGE_IMPORT_DESCRIPTOR pImportDesc = 0;
	PIMAGE_IMPORT_BY_NAME pImportByName = 0;
	PIMAGE_THUNK_DATA pThunk = 0;
	ImportModuleThunk * importModuleThunk = 0;
	ImportThunk * importThunk = 0;

	size_t stringLength = 0;
	DWORD_PTR lastRVA = 0;

	BYTE * sectionData = listPeSection[importSectionIndex].data;
	DWORD offset = 0;

	pImportDescriptor = (PIMAGE_IMPORT_DESCRIPTOR)(sectionData);

	//skip the IMAGE_IMPORT_DESCRIPTOR
	offset += (DWORD)(numberOfImportDescriptors * sizeof(IMAGE_IMPORT_DESCRIPTOR));

	for ( mapIt = moduleList.begin() ; mapIt != moduleList.end(); mapIt++ )
	{
	importModuleThunk = &((*mapIt).second);

	stringLength = addImportDescriptor(importModuleThunk, offset);

	#ifdef DEBUG_COMMENTS
	Scylla::debugLog.log(L"fillImportSection :: importDesc.Name %X", pImportDescriptor->Name);
	#endif

	offset += (DWORD)stringLength; //stringLength has null termination char

	pImportByName = (PIMAGE_IMPORT_BY_NAME)((DWORD_PTR)sectionData + offset);

	//pThunk = (PIMAGE_THUNK_DATA)(getMemoryPointerFromRVA(importModuleThunk->firstThunk));

	lastRVA = importModuleThunk->firstThunk - sizeof(DWORD_PTR);

	for ( mapIt2 = (mapIt).second.thunkList.begin() ; mapIt2 != (mapIt).second.thunkList.end(); mapIt2++ )
	{
	importThunk = &((*mapIt2).second);

	pThunk = (PIMAGE_THUNK_DATA)(getMemoryPointerFromRVA(importThunk->rva));

	//check wrong iat pointer
	if (!pThunk)
	{
	#ifdef DEBUG_COMMENTS
	Scylla::debugLog.log(L"fillImportSection :: Failed to get pThunk RVA: %X", importThunk->rva);
	#endif
	return 0;
	}

	if ((lastRVA + sizeof(DWORD_PTR)) != importThunk->rva)
	{
	//add additional import desc
	addSpecialImportDescriptor(importThunk->rva);
	}
	lastRVA = importThunk->rva;

	#ifdef DEBUG_COMMENTS
	Scylla::debugLog.log(L"fillImportSection :: importThunk %X pThunk %X pImportByName %X offset %X", importThunk,pThunk,pImportByName,offset);
	#endif
	stringLength = addImportToImportTable(importThunk, pThunk, pImportByName, offset);

	offset += (DWORD)stringLength; //is 0 bei import by ordinal
	pImportByName = (PIMAGE_IMPORT_BY_NAME)((DWORD_PTR)pImportByName + stringLength);
	}

	pImportDescriptor++;
	}

	return offset;
	}

	size_t ImportRebuilder::addImportDescriptor(ImportModuleThunk * pImportModule, DWORD sectionOffset)
	{
	char dllName[MAX_PATH];

	StringConversion::ToASCII(pImportModule->moduleName, dllName, _countof(dllName));
	size_t stringLength = strlen(dllName) + 1;

	/*
	Warning: stringLength MUST include null termination char
	*/

	memcpy((listPeSection[importSectionIndex].data + sectionOffset), dllName, stringLength); //copy module name to section

	pImportDescriptor->FirstThunk = (DWORD)pImportModule->firstThunk;
	pImportDescriptor->Name = (DWORD)convertOffsetToRVAVector(listPeSection[importSectionIndex].sectionHeader.PointerToRawData + sectionOffset);

	return stringLength;
	}

	void ImportRebuilder::addSpecialImportDescriptor(DWORD_PTR rvaFirstThunk)
	{
	PIMAGE_IMPORT_DESCRIPTOR oldID = pImportDescriptor;
	pImportDescriptor++;

	pImportDescriptor->FirstThunk = (DWORD)rvaFirstThunk;
	pImportDescriptor->Name = oldID->Name;
	}

	void ImportRebuilder::calculateImportSizes(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	std::map<DWORD_PTR, ImportModuleThunk>::iterator mapIt;
	std::map<DWORD_PTR, ImportThunk>::iterator mapIt2;
	DWORD_PTR lastRVA = 0;

	numberOfImportDescriptors = 0;
	sizeOfImportSection = 0;
	sizeOfApiAndModuleNames = 0;

	numberOfImportDescriptors = moduleList.size() + 1; //last is zero'd

	for ( mapIt = moduleList.begin() ; mapIt != moduleList.end(); mapIt++ )
	{
	lastRVA = (*mapIt).second.firstThunk - sizeof(DWORD_PTR);

	sizeOfApiAndModuleNames += (DWORD)(wcslen((*mapIt).second.moduleName) + 1);

	for ( mapIt2 = (mapIt).second.thunkList.begin() ; mapIt2 != (mapIt).second.thunkList.end(); mapIt2++ )
	{
	if ((lastRVA + sizeof(DWORD_PTR)) != (*mapIt2).second.rva)
	{
	numberOfImportDescriptors++; //add additional import desc
	}

	if((*mapIt2).second.name[0] != '\0')
	{
	sizeOfApiAndModuleNames += sizeof(WORD); //Hint from IMAGE_IMPORT_BY_NAME
	sizeOfApiAndModuleNames += (DWORD)(strlen((*mapIt2).second.name) + 1);
	}

	lastRVA = (*mapIt2).second.rva;
	}
	}

	sizeOfImportSection = sizeOfApiAndModuleNames + (numberOfImportDescriptors * sizeof(IMAGE_IMPORT_DESCRIPTOR));
	}

	size_t ImportRebuilder::addImportToImportTable( ImportThunk * pImport, PIMAGE_THUNK_DATA pThunk, PIMAGE_IMPORT_BY_NAME pImportByName, DWORD sectionOffset)
	{
	size_t stringLength = 0;

	if(pImport->name[0] == '\0')
	{
	pThunk->u1.AddressOfData = (IMAGE_ORDINAL(pImport->ordinal) \| IMAGE_ORDINAL_FLAG);
	}
	else
	{
	pImportByName->Hint = pImport->hint;

	stringLength = strlen(pImport->name) + 1;
	memcpy(pImportByName->Name, pImport->name, stringLength);

	pThunk->u1.AddressOfData = convertOffsetToRVAVector(listPeSection[importSectionIndex].sectionHeader.PointerToRawData + sectionOffset);

	if (!pThunk->u1.AddressOfData)
	{
	#ifdef DEBUG_COMMENTS
	Scylla::debugLog.log(L"addImportToImportTable :: failed to get AddressOfData %X %X", listPeSection[importSectionIndex].sectionHeader.PointerToRawData, sectionOffset);
	#endif
	}

	//next import should be nulled
	pThunk++;
	pThunk->u1.AddressOfData = 0;

	#ifdef DEBUG_COMMENTS
	Scylla::debugLog.log(L"addImportToImportTable :: pThunk->u1.AddressOfData %X %X %X", pThunk->u1.AddressOfData, pThunk, listPeSection[importSectionIndex].sectionHeader.PointerToRawData + sectionOffset);
	#endif
	stringLength += sizeof(WORD);
	}

	return stringLength;
	}

	BYTE * ImportRebuilder::getMemoryPointerFromRVA(DWORD_PTR dwRVA)
	{
	DWORD_PTR offset = convertRVAToOffsetVector(dwRVA);

	for (size_t i = 0; i < listPeSection.size(); i++)
	{
	if ((listPeSection[i].sectionHeader.PointerToRawData <= offset) && ((listPeSection[i].sectionHeader.PointerToRawData + listPeSection[i].sectionHeader.SizeOfRawData) > offset))
	{
	return (BYTE *)((DWORD_PTR)listPeSection[i].data + (offset - listPeSection[i].sectionHeader.PointerToRawData));
	}
	}

	return 0;
	}

File Metadata

Mime Type: text/x-c
Expires: Tue, Feb 10, 11:29 PM (1 d, 8 h)
Storage Engine: local-disk
Storage Format: Raw Data
Storage Handle: 83/b4/c9fab536339efabbbf6d8ba09f9a

ImportRebuilder.cppNo OneTemporaryActions

ImportRebuilder.cppView Options

File Metadata

Event Timeline

ImportRebuilder.cpp
No OneTemporary
Actions

ImportRebuilder.cpp
View Options