No OneTemporary
Actions

Size

51 KB

Subscribers

None

View Options

	diff --git a/Scylla/ImportRebuild.cpp b/Scylla/ImportRebuild.cpp
	index 4129731..f396b62 100644
	--- a/Scylla/ImportRebuild.cpp
	+++ b/Scylla/ImportRebuild.cpp
	@@ -1,670 +1,698 @@
	#include "ImportRebuild.h"

	#include "Logger.h"
	#include "ConfigurationHolder.h"
	-//#define DEBUG_COMMENTS
	+
	+#define DEBUG_COMMENTS


	bool ImportRebuild::splitTargetFile()
	{
	PIMAGE_SECTION_HEADER pSecHeader = 0;
	WORD i = 0;
	BYTE * data = 0;
	DWORD alignment = 0;
	DWORD dwSize = 0;

	pDosHeader = new IMAGE_DOS_HEADER;
	CopyMemory(pDosHeader, imageData, sizeof(IMAGE_DOS_HEADER));

	if (pDosHeader->e_magic != IMAGE_DOS_SIGNATURE)
	{
	return false;
	}

	pNTHeader = new IMAGE_NT_HEADERS;
	CopyMemory(pNTHeader, (PVOID)((DWORD_PTR)imageData + pDosHeader->e_lfanew), sizeof(IMAGE_NT_HEADERS));

	if (pNTHeader->Signature != IMAGE_NT_SIGNATURE)
	{
	return false;
	}

	if (pDosHeader->e_lfanew > sizeof(IMAGE_DOS_HEADER))
	{
	dwSize = pDosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER);
	pDosStub = new BYTE[dwSize];
	CopyMemory(pDosStub, (PVOID)((DWORD_PTR)imageData + sizeof(IMAGE_DOS_HEADER)), dwSize);
	}
	else
	{
	pDosStub = 0;
	}

	pSecHeader = IMAGE_FIRST_SECTION((PIMAGE_NT_HEADERS)((DWORD_PTR)imageData + pDosHeader->e_lfanew));

	for (i = 0; i < pNTHeader->FileHeader.NumberOfSections; i++)
	{
	dwSize = pSecHeader->SizeOfRawData;

	if (dwSize > 300000000)
	{
	dwSize = 300000000;
	}

	//TODO better use section alignment because it is better?
	alignment = alignValue(dwSize, pNTHeader->OptionalHeader.SectionAlignment);
	data = new BYTE[alignment];

	ZeroMemory(data, alignment);
	CopyMemory(data, (PVOID)((DWORD_PTR)imageData + pSecHeader->PointerToRawData), dwSize);

	vecSectionData.push_back(data);
	vecSectionHeaders.push_back(*pSecHeader);

	pSecHeader++;
	}

	delete [] imageData;
	imageData = 0;

	return true;
	}

	bool ImportRebuild::alignSectionHeaders()
	{
	for (WORD i = 0; i < vecSectionHeaders.size(); i++)
	{
	vecSectionHeaders[i].VirtualAddress = alignValue(vecSectionHeaders[i].VirtualAddress, pNTHeader->OptionalHeader.SectionAlignment);
	vecSectionHeaders[i].Misc.VirtualSize = alignValue(vecSectionHeaders[i].Misc.VirtualSize, pNTHeader->OptionalHeader.SectionAlignment);

	vecSectionHeaders[i].PointerToRawData = alignValue(vecSectionHeaders[i].PointerToRawData, pNTHeader->OptionalHeader.FileAlignment);
	vecSectionHeaders[i].SizeOfRawData = alignValue(vecSectionHeaders[i].SizeOfRawData, pNTHeader->OptionalHeader.FileAlignment);
	}

	return true;
	}

	bool ImportRebuild::saveNewFile(const WCHAR * filepath)
	{
	DWORD fileOffset = 0;
	DWORD dwWriteSize = 0;
	size_t i = 0;

	if (vecSectionHeaders.size() != vecSectionData.size())
	{
	return false;
	}

	HANDLE hFile = CreateFile(filepath, GENERIC_WRITE, FILE_SHARE_WRITE \| FILE_SHARE_READ, 0,CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);

	if(hFile == INVALID_HANDLE_VALUE)
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog("saveNewFile :: INVALID_HANDLE_VALUE %u\r\n",GetLastError());
	#endif

	return false;
	}

	//alignSectionHeaders();
	updatePeHeader();

	fileOffset = 0;
	dwWriteSize = sizeof(IMAGE_DOS_HEADER);
	ProcessAccessHelp::writeMemoryToFile(hFile, fileOffset, dwWriteSize, pDosHeader);

	fileOffset += dwWriteSize;
	dwWriteSize = pDosHeader->e_lfanew - sizeof(IMAGE_DOS_HEADER);
	ProcessAccessHelp::writeMemoryToFile(hFile, fileOffset, dwWriteSize, pDosStub);

	fileOffset += dwWriteSize;
	dwWriteSize = sizeof(IMAGE_NT_HEADERS);
	ProcessAccessHelp::writeMemoryToFile(hFile, fileOffset, dwWriteSize, pNTHeader);

	fileOffset += dwWriteSize;
	dwWriteSize = sizeof(IMAGE_SECTION_HEADER);

	for (i = 0; i < vecSectionHeaders.size(); i++)
	{
	if (!ProcessAccessHelp::writeMemoryToFile(hFile, fileOffset, dwWriteSize, &vecSectionHeaders[i]))
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog(TEXT("saveNewFile :: writeMemoryToFile failed offset %X size %X\r\n"),fileOffset,dwWriteSize);
	#endif
	CloseHandle(hFile);
	return false;
	}
	fileOffset += dwWriteSize;
	}

	for (i = 0; i < vecSectionHeaders.size(); i++)
	{
	dwWriteSize = vecSectionHeaders[i].PointerToRawData - fileOffset;

	if (dwWriteSize)
	{
	if (!writeZeroMemoryToFile(hFile, fileOffset, dwWriteSize))
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog(TEXT("saveNewFile :: writeZeroMemoryToFile failed offset %X size %X\r\n"),fileOffset,dwWriteSize);
	#endif
	CloseHandle(hFile);
	return false;
	}
	fileOffset += dwWriteSize;
	}

	dwWriteSize = vecSectionHeaders[i].SizeOfRawData;

	ProcessAccessHelp::writeMemoryToFile(hFile, fileOffset, dwWriteSize, vecSectionData[i]);
	fileOffset += dwWriteSize;
	}


	CloseHandle(hFile);

	return true;
	}

	bool ImportRebuild::writeZeroMemoryToFile(HANDLE hFile, DWORD fileOffset, DWORD size)
	{
	bool retValue = false;
	PVOID zeromemory = calloc(size, 1);

	if (zeromemory)
	{
	retValue = ProcessAccessHelp::writeMemoryToFile(hFile, fileOffset, size, zeromemory);
	free(zeromemory);
	}
	else
	{
	retValue = false;
	}

	return retValue;
	}

	void ImportRebuild::clearAllData()
	{
	if (pDosStub)
	{
	delete [] pDosStub;
	pDosStub = 0;
	}

	if (imageData)
	{
	delete [] imageData;
	imageData = 0;
	}

	delete pDosHeader;
	pDosHeader = 0;

	delete pNTHeader;
	pNTHeader = 0;

	vecSectionHeaders.clear();

	for (size_t i = 0; i < vecSectionData.size(); i++)
	{
	delete [] vecSectionData[i];
	}

	vecSectionData.clear();
	}

	bool ImportRebuild::addNewSection(char * sectionName, DWORD sectionSize, BYTE * sectionData)
	{
	BYTE * newBuffer = 0;
	IMAGE_SECTION_HEADER pNewSection = {0};
	size_t lastSectionIndex = vecSectionHeaders.size() - 1;
	size_t nameLength = strlen(sectionName);

	if (nameLength > IMAGE_SIZEOF_SHORT_NAME)
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog(TEXT("addNewSection :: sectionname is too long %d\r\n"),nameLength);
	#endif
	return false;
	}

	memcpy_s(pNewSection.Name, IMAGE_SIZEOF_SHORT_NAME, sectionName, nameLength);

	pNewSection.SizeOfRawData = alignValue(sectionSize, pNTHeader->OptionalHeader.FileAlignment);
	pNewSection.Misc.VirtualSize = alignValue(sectionSize, pNTHeader->OptionalHeader.SectionAlignment);

	pNewSection.PointerToRawData = alignValue(vecSectionHeaders[lastSectionIndex].PointerToRawData + vecSectionHeaders[lastSectionIndex].SizeOfRawData, pNTHeader->OptionalHeader.FileAlignment);
	pNewSection.VirtualAddress = alignValue(vecSectionHeaders[lastSectionIndex].VirtualAddress + vecSectionHeaders[lastSectionIndex].Misc.VirtualSize, pNTHeader->OptionalHeader.SectionAlignment);

	pNewSection.Characteristics = IMAGE_SCN_MEM_EXECUTE\|IMAGE_SCN_MEM_READ\|IMAGE_SCN_MEM_WRITE\|IMAGE_SCN_CNT_CODE\|IMAGE_SCN_CNT_INITIALIZED_DATA;

	vecSectionHeaders.push_back(pNewSection);

	if ( (sectionSize != pNewSection.SizeOfRawData) \|\| (sectionData == 0) )
	{
	newBuffer = new BYTE[pNewSection.SizeOfRawData];
	ZeroMemory(newBuffer, pNewSection.SizeOfRawData);

	if (sectionData)
	{
	CopyMemory(newBuffer, sectionData, sectionSize);
	}

	}
	else
	{
	newBuffer = sectionData;
	}

	vecSectionData.push_back(newBuffer);

	return true;
	}

	bool ImportRebuild::loadTargetFile(const WCHAR * filepath)
	{
	HANDLE hTargetFile = INVALID_HANDLE_VALUE;
	DWORD fileSize = 0;
	bool retValue = false;

	hTargetFile = CreateFile(filepath, GENERIC_READ\|GENERIC_WRITE, FILE_SHARE_READ\|FILE_SHARE_WRITE, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);

	if(hTargetFile == INVALID_HANDLE_VALUE)
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog("loadTargetFile :: INVALID_HANDLE_VALUE %u\r\n",GetLastError());
	#endif

	return false;
	}

	fileSize = (DWORD)ProcessAccessHelp::getFileSize(hTargetFile);

	if (!fileSize)
	{
	CloseHandle(hTargetFile);
	hTargetFile = 0;
	return false;
	}

	imageData = new BYTE[fileSize];

	if (!imageData)
	{
	retValue = false;
	}
	else
	{
	retValue = ProcessAccessHelp::readMemoryFromFile(hTargetFile, 0, fileSize, imageData);
	}

	CloseHandle(hTargetFile);
	hTargetFile = 0;

	return retValue;
	}

	DWORD ImportRebuild::alignValue(DWORD badValue, DWORD alignTo)
	{
	return (((badValue + alignTo - 1) / alignTo) * alignTo);
	}

	DWORD ImportRebuild::convertRVAToOffsetVector(DWORD dwRVA)
	{
	for (size_t i = 0; i < vecSectionHeaders.size(); i++)
	{
	if ((vecSectionHeaders[i].VirtualAddress <= dwRVA) && ((vecSectionHeaders[i].VirtualAddress + vecSectionHeaders[i].Misc.VirtualSize) > dwRVA))
	{
	return ((dwRVA - vecSectionHeaders[i].VirtualAddress) + vecSectionHeaders[i].PointerToRawData);
	}
	}

	return 0;
	}

	DWORD ImportRebuild::convertRVAToOffset(DWORD dwRVA)
	{
	PIMAGE_SECTION_HEADER pSectionHeader = IMAGE_FIRST_SECTION(pNTHeader);

	for (WORD i = 0; i < pNTHeader->FileHeader.NumberOfSections; i++)
	{
	if ((pSectionHeader->VirtualAddress <= dwRVA) && ((pSectionHeader->VirtualAddress + pSectionHeader->Misc.VirtualSize) > dwRVA))
	{
	return ((dwRVA - pSectionHeader->VirtualAddress) + pSectionHeader->PointerToRawData);
	}
	pSectionHeader++;
	}

	return 0;
	}

	DWORD_PTR ImportRebuild::convertOffsetToRVAVector(DWORD dwOffset)
	{
	for (size_t i = 0; i < vecSectionHeaders.size(); i++)
	{
	if ((vecSectionHeaders[i].PointerToRawData <= dwOffset) && ((vecSectionHeaders[i].PointerToRawData + vecSectionHeaders[i].SizeOfRawData) > dwOffset))
	{
	return ((dwOffset - vecSectionHeaders[i].PointerToRawData) + vecSectionHeaders[i].VirtualAddress);
	}
	}

	return 0;
	}

	DWORD ImportRebuild::convertOffsetToRVA(DWORD dwOffset)
	{
	PIMAGE_SECTION_HEADER pSectionHeader = IMAGE_FIRST_SECTION(pNTHeader);

	for (WORD i = 0; i < pNTHeader->FileHeader.NumberOfSections; i++)
	{
	if ((pSectionHeader->PointerToRawData <= dwOffset) && ((pSectionHeader->PointerToRawData + pSectionHeader->SizeOfRawData) > dwOffset))
	{
	return ((dwOffset - pSectionHeader->PointerToRawData) + pSectionHeader->VirtualAddress);
	}
	pSectionHeader++;
	}

	return 0;
	}

	void ImportRebuild::updatePeHeader()
	{
	size_t lastSectionIndex = vecSectionHeaders.size() - 1;

	pNTHeader->FileHeader.NumberOfSections = (WORD)(lastSectionIndex + 1);
	pNTHeader->OptionalHeader.SizeOfImage = vecSectionHeaders[lastSectionIndex].VirtualAddress + vecSectionHeaders[lastSectionIndex].Misc.VirtualSize;

	pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].VirtualAddress = 0;
	pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].Size = 0;

	if (pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress)
	{
	for (size_t i = 0; i < vecSectionHeaders.size(); i++)
	{
	if ((vecSectionHeaders[i].VirtualAddress <= pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress) && ((vecSectionHeaders[i].VirtualAddress + vecSectionHeaders[i].Misc.VirtualSize) > pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress))
	{
	//section must be read and writeable
	vecSectionHeaders[i].Characteristics \|= IMAGE_SCN_MEM_READ\|IMAGE_SCN_MEM_WRITE;
	}
	}

	pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress = 0;
	pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size = 0;
	}


	pNTHeader->OptionalHeader.NumberOfRvaAndSizes = 0x10;

	pNTHeader->OptionalHeader.SizeOfHeaders = alignValue(pDosHeader->e_lfanew + sizeof(DWORD) + sizeof(IMAGE_FILE_HEADER) + pNTHeader->FileHeader.SizeOfOptionalHeader + (pNTHeader->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER)), pNTHeader->OptionalHeader.FileAlignment);

	}



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

	importSectionIndex = vecSectionHeaders.size() - 1;

	DWORD dwSize = fillImportSection(moduleList);

	+ if (!dwSize)
	+ {
	+ return false;
	+ }
	+
	setFlagToIATSection((*moduleList.begin()).second.firstThunk);

	pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress = vecSectionHeaders[importSectionIndex].VirtualAddress;
	pNTHeader->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size = (DWORD)(numberOfImportDescriptors * sizeof(IMAGE_IMPORT_DESCRIPTOR));
	return true;
	}

	bool ImportRebuild::createNewImportSection(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	char sectionName[9] = {0};
	size_t i = 0;

	//DWORD sectionSize = calculateMinSize(moduleList);
	calculateImportSizes(moduleList);

	if (wcslen(ConfigurationHolder::getConfigObject(IAT_SECTION_NAME)->valueString) > IMAGE_SIZEOF_SHORT_NAME)
	{
	strcpy_s(sectionName, sizeof(sectionName), ".SCY");
	}
	else
	{
	wcstombs_s(&i, sectionName, sizeof(sectionName), ConfigurationHolder::getConfigObject(IAT_SECTION_NAME)->valueString, _TRUNCATE);
	}


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

	/*DWORD ImportRebuild::calculateMinSize(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
	{
	DWORD dwSize = 0;
	std::map<DWORD_PTR, ImportModuleThunk>::iterator mapIt;
	std::map<DWORD_PTR, ImportThunk>::iterator mapIt2;

	dwSize = (DWORD)((moduleList.size() + 1) * sizeof(IMAGE_IMPORT_DESCRIPTOR)); //last is zero'ed

	for ( mapIt = moduleList.begin() ; mapIt != moduleList.end(); mapIt++ )
	{

	//dwSize += (DWORD)((*mapIt).second.thunkList.size() + sizeof(IMAGE_IMPORT_BY_NAME));
	dwSize += (DWORD)(wcslen((*mapIt).second.moduleName) + 1);

	for ( mapIt2 = (mapIt).second.thunkList.begin() ; mapIt2 != (mapIt).second.thunkList.end(); mapIt2++ )
	{
	if((*mapIt2).second.name[0] != '\0')
	{
	dwSize += sizeof(IMAGE_IMPORT_BY_NAME);
	dwSize += (DWORD)strlen((*mapIt2).second.name);
	}
	}
	}

	return dwSize;
	}*/

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

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

	return 0;
	}

	DWORD ImportRebuild::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 = vecSectionData[importSectionIndex];
	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
	Logger::debugLog("fillImportSection :: importDesc.Name %X\r\n", 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
	+ Logger::debugLog(TEXT("fillImportSection :: Failed to get pThunk RVA: %X\n"), 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
	+ Logger::debugLog(TEXT("fillImportSection :: importThunk %X pThunk %X pImportByName %X offset %X\n"), 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;
	}

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

	if (loadTargetFile(targetFilePath))
	{
	splitTargetFile();

	- buildNewImportTable(moduleList);
	+ retValue = buildNewImportTable(moduleList);

	- retValue = saveNewFile(newFilePath);
	+ if (retValue)
	+ {
	+ retValue = saveNewFile(newFilePath);
	+ }

	clearAllData();

	return retValue;
	}
	else
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog(TEXT("rebuildImportTable ::Failed to load target %s\n"), targetFilePath);
	#endif
	return false;
	}
	}

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

	size_t ImportRebuild::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(vecSectionHeaders[importSectionIndex].PointerToRawData + sectionOffset);

	+ if (!pThunk->u1.AddressOfData)
	+ {
	+#ifdef DEBUG_COMMENTS
	+ Logger::debugLog("addImportToImportTable :: failed to get AddressOfData %X %X\n",vecSectionHeaders[importSectionIndex].PointerToRawData, sectionOffset);
	+#endif
	+ }
	+
	#ifdef DEBUG_COMMENTS
	- Logger::debugLog("pThunk->u1.AddressOfData %X %X %X\n",pThunk->u1.AddressOfData, pThunk, vecSectionHeaders[importSectionIndex].PointerToRawData + offset);
	+ Logger::debugLog("addImportToImportTable :: pThunk->u1.AddressOfData %X %X %X\n",pThunk->u1.AddressOfData, pThunk, vecSectionHeaders[importSectionIndex].PointerToRawData + sectionOffset);
	#endif
	stringLength += sizeof(WORD);
	}

	return stringLength;
	}

	size_t ImportRebuild::addImportDescriptor(ImportModuleThunk * pImportModule, DWORD sectionOffset)
	{
	char dllName[MAX_PATH];
	size_t stringLength = 0;

	wcstombs_s(&stringLength, dllName, (size_t)MAX_PATH, pImportModule->moduleName, (size_t)MAX_PATH);

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

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

	return stringLength;
	}

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

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

	void ImportRebuild::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));
	}
	\ No newline at end of file
	diff --git a/Scylla/MainGui.cpp b/Scylla/MainGui.cpp
	index 22b4ddc..e6f01cc 100644
	--- a/Scylla/MainGui.cpp
	+++ b/Scylla/MainGui.cpp
	@@ -1,1238 +1,1241 @@
	#include "MainGui.h"

	#include <atldlgs.h> // WTL common dialogs

	#include "definitions.h"
	#include "PluginLoader.h"
	#include "ConfigurationHolder.h"
	#include "PeDump.h"
	#include "PeRebuild.h"
	#include "DllInjectionPlugin.h"
	#include "DisassemblerGui.h"
	#include "PickApiGui.h"
	#include "NativeWinApi.h"
	#include "ImportRebuild.h"
	#include "SystemInformation.h"
	#include "AboutGui.h"
	#include "OptionsGui.h"

	extern CAppModule _Module; // o_O

	const WCHAR MainGui::filterExe[] = L"Executable (.exe)\0.exe\0All files\0.\0";
	const WCHAR MainGui::filterDll[] = L"Dynamic Link Library (.dll)\0.dll\0All files\0.\0";
	const WCHAR MainGui::filterExeDll[] = L"Executable (.exe)\0.exe\0Dynamic Link Library (.dll)\0.dll\0All files\0.\0";
	const WCHAR MainGui::filterTxt[] = L"Text file (.txt)\0.txt\0All files\0.\0";

	MainGui::MainGui() : selectedProcess(0), importsHandling(TreeImports), TreeImportsSubclass(this, IDC_TREE_IMPORTS)
	{
	Logger::getDebugLogFilePath();
	ConfigurationHolder::loadConfiguration();
	PluginLoader::findAllPlugins();
	NativeWinApi::initialize();
	SystemInformation::getSystemInformation();

	if(ConfigurationHolder::getConfigObject(DEBUG_PRIVILEGE)->isTrue())
	{
	processLister.setDebugPrivileges();
	}

	processAccessHelp.getProcessModules(GetCurrentProcessId(), processAccessHelp.ownModuleList);

	hIcon.LoadIcon(IDI_ICON_SCYLLA);
	hMenuImports.LoadMenu(IDR_MENU_IMPORTS);
	hMenuLog.LoadMenu(IDR_MENU_LOG);
	accelerators.LoadAccelerators(IDR_ACCELERATOR_MAIN);

	hIconCheck.LoadIcon(IDI_ICON_CHECK, 16, 16);
	hIconWarning.LoadIcon(IDI_ICON_WARNING, 16, 16);
	hIconError.LoadIcon(IDI_ICON_ERROR, 16, 16);

	if(hMenuImports)
	{
	appendPluginListToMenu(hMenuImports.GetSubMenu(0));
	}
	}

	BOOL MainGui::PreTranslateMessage(MSG* pMsg)
	{
	if(accelerators.TranslateAccelerator(m_hWnd, pMsg))
	{
	return TRUE;
	}
	else if(IsDialogMessage(pMsg))
	{
	return TRUE;
	}

	return FALSE;
	}

	BOOL MainGui::OnInitDialog(CWindow wndFocus, LPARAM lInitParam)
	{
	if (SystemInformation::currenOS == UNKNOWN_OS)
	{
	if(IDCANCEL == MessageBox(L"Operating System is not supported\r\nContinue anyway?", L"Scylla", MB_ICONWARNING \| MB_OKCANCEL))
	{
	SendMessage(WM_CLOSE);
	return FALSE;
	}
	}

	CMessageLoop* pLoop = _Module.GetMessageLoop();
	pLoop->AddMessageFilter(this);

	setupStatusBar();
	fillStatusBar();

	DoDataExchange(); // attach controls
	DlgResize_Init(true, true);

	EditOEPAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH);
	EditIATAddress.LimitText(MAX_HEX_VALUE_EDIT_LENGTH);
	EditIATSize.LimitText(MAX_HEX_VALUE_EDIT_LENGTH);

	appendPluginListToMenu(CMenuHandle(GetMenu()).GetSubMenu(MenuImportsOffsetTrace));

	enableDialogControls(FALSE);

	setIconAndDialogCaption();

	return TRUE;
	}

	void MainGui::OnDestroy()
	{
	PostQuitMessage(0);
	}

	void MainGui::OnSize(UINT nType, CSize size)
	{
	StatusBar.SendMessage(WM_SIZE);
	SetMsgHandled(FALSE);
	}

	void MainGui::OnLButtonDown(UINT nFlags, CPoint point)
	{
	SetMsgHandled(FALSE);
	}

	void MainGui::OnContextMenu(CWindow wnd, CPoint point)
	{
	// point = -1, -1 for keyboard invoked shortcut!
	switch(wnd.GetDlgCtrlID())
	{
	case IDC_TREE_IMPORTS:
	DisplayContextMenuImports(wnd, point);
	return;
	case IDC_LIST_LOG:
	DisplayContextMenuLog(wnd, point);
	return;
	//default: // wnd == m_hWnd?
	// DisplayContextMenu(wnd, point);
	// return;
	}

	SetMsgHandled(FALSE);
	}

	void MainGui::OnCommand(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// Make sure it's a menu
	if(uNotifyCode == 0 && !wndCtl.IsWindow())
	{
	if ((nID >= PLUGIN_MENU_BASE_ID) && (nID <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID)))
	{
	pluginActionHandler(nID);
	return;
	}
	}
	SetMsgHandled(FALSE);
	}

	LRESULT MainGui::OnTreeImportsClick(const NMHDR* pnmh)
	{
	SetMsgHandled(FALSE);
	return 0;
	}

	LRESULT MainGui::OnTreeImportsDoubleClick(const NMHDR* pnmh)
	{
	if(TreeImports.GetCount() < 1)
	return 0;

	// Get item under cursor
	CPoint client = GetMessagePos();
	TreeImports.ScreenToClient(&client);
	UINT flags;
	CTreeItem over = TreeImports.HitTest(client, &flags);
	CTreeItem parent;
	if(over)
	{
	if(!(flags & TVHT_ONITEM))
	{
	over = NULL;
	}
	else
	{
	parent = over.GetParent();
	}
	}

	if(!over.IsNull() && !parent.IsNull())
	{
	pickApiActionHandler(over);
	}

	return 0;
	}

	LRESULT MainGui::OnTreeImportsRightClick(const NMHDR* pnmh)
	{
	SetMsgHandled(FALSE);
	return 0;
	}

	LRESULT MainGui::OnTreeImportsRightDoubleClick(const NMHDR* pnmh)
	{
	SetMsgHandled(FALSE);
	return 0;
	}

	LRESULT MainGui::OnTreeImportsKeyDown(const NMHDR* pnmh)
	{
	const NMTVKEYDOWN * tkd = (NMTVKEYDOWN *)pnmh;
	switch(tkd->wVKey)
	{
	case VK_RETURN:
	{
	CTreeItem selected = TreeImports.GetFirstSelectedItem();
	if(!selected.IsNull() && !selected.GetParent().IsNull())
	{
	pickApiActionHandler(selected);
	}
	}
	return 1;
	case VK_DELETE:
	deleteSelectedImportsActionHandler();
	return 1;
	}

	SetMsgHandled(FALSE);
	return 0;
	}

	void MainGui::deleteSelectedImportsActionHandler()
	{
	CTreeItem selected = TreeImports.GetFirstSelectedItem();
	while(!selected.IsNull())
	{
	if(importsHandling.isModule(selected))
	{
	importsHandling.cutModule(selected);
	}
	else
	{
	importsHandling.cutImport(selected);
	}
	selected = TreeImports.GetNextSelectedItem(selected);
	}
	fillStatusBar();
	}

	void MainGui::invalidateSelectedImportsActionHandler()
	{
	CTreeItem selected = TreeImports.GetFirstSelectedItem();
	while(!selected.IsNull())
	{
	if(importsHandling.isImport(selected))
	{
	importsHandling.invalidateImport(selected);
	}
	selected = TreeImports.GetNextSelectedItem(selected);
	}
	fillStatusBar();
	}

	UINT MainGui::OnTreeImportsSubclassGetDlgCode(const MSG * lpMsg)
	{
	if(lpMsg)
	{
	switch(lpMsg->wParam)
	{
	case VK_RETURN:
	return DLGC_WANTMESSAGE;
	}
	}

	SetMsgHandled(FALSE);
	return 0;
	}

	void MainGui::OnTreeImportsSubclassChar(UINT nChar, UINT nRepCnt, UINT nFlags)
	{
	switch(nChar)
	{
	case VK_RETURN:
	break;
	default:
	SetMsgHandled(FALSE);
	break;
	}
	}

	void MainGui::OnProcessListDrop(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	fillProcessListComboBox(ComboProcessList);
	}

	void MainGui::OnProcessListSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	processSelectedActionHandler(ComboProcessList.GetCurSel());
	}

	void MainGui::OnPickDLL(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	pickDllActionHandler();
	}

	void MainGui::OnOptions(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	optionsActionHandler();
	}

	void MainGui::OnDump(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	dumpActionHandler();
	}

	void MainGui::OnFixDump(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	dumpFixActionHandler();
	}

	void MainGui::OnPERebuild(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	peRebuildActionHandler();
	}

	void MainGui::OnDLLInject(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	dllInjectActionHandler();
	}

	void MainGui::OnIATAutoSearch(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	iatAutosearchActionHandler();
	}

	void MainGui::OnGetImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	getImportsActionHandler();
	}

	void MainGui::OnInvalidImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	showInvalidImportsActionHandler();
	}

	void MainGui::OnSuspectImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	showSuspectImportsActionHandler();
	}

	void MainGui::OnClearImports(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	clearImportsActionHandler();
	}

	void MainGui::OnInvalidateSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	invalidateSelectedImportsActionHandler();
	}

	void MainGui::OnCutSelected(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	deleteSelectedImportsActionHandler();
	}

	void MainGui::OnSaveTree(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// TODO
	}

	void MainGui::OnLoadTree(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// TODO
	}

	void MainGui::OnAutotrace(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	// TODO
	}

	void MainGui::OnExit(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	DestroyWindow();
	}

	void MainGui::OnAbout(UINT uNotifyCode, int nID, CWindow wndCtl)
	{
	showAboutDialog();
	}

	void MainGui::setupStatusBar()
	{
	StatusBar.Create(m_hWnd, NULL, L"", WS_CHILD \| WS_VISIBLE \| WS_CLIPCHILDREN \| WS_CLIPSIBLINGS \| SBARS_TOOLTIPS, NULL, IDC_STATUS_BAR);

	CRect rcMain, rcStatus;
	GetClientRect(&rcMain);
	StatusBar.GetWindowRect(&rcStatus);

	const int PARTS = 4;
	int widths[PARTS];

	widths[PART_COUNT] = rcMain.Width() / 5;
	widths[PART_INVALID] = widths[PART_COUNT] + rcMain.Width() / 5;
	widths[PART_IMAGEBASE] = widths[PART_INVALID] + rcMain.Width() / 3;
	widths[PART_MODULE] = -1;

	StatusBar.SetParts(PARTS, widths);

	ResizeClient(rcMain.Width(), rcMain.Height() + rcStatus.Height(), FALSE);
	}

	void MainGui::fillStatusBar()
	{
	// Rewrite ImportsHandling so we get these easily
	unsigned int totalImports = importsHandling.thunkCount();
	unsigned int invalidImports = importsHandling.invalidThunkCount();

	// \t = center, \t\t = right-align
	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("\tImports: %u"), totalImports);
	StatusBar.SetText(PART_COUNT, stringBuffer);

	if(invalidImports > 0)
	{
	StatusBar.SetIcon(PART_INVALID, hIconError);
	}
	else
	{
	StatusBar.SetIcon(PART_INVALID, hIconCheck);
	}

	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("\tInvalid: %u"), invalidImports);
	StatusBar.SetText(PART_INVALID, stringBuffer);

	if(selectedProcess)
	{
	DWORD_PTR imageBase = 0;
	const WCHAR * fileName = 0;

	if(processAccessHelp.selectedModule)
	{
	imageBase = processAccessHelp.selectedModule->modBaseAddr;
	fileName = processAccessHelp.selectedModule->getFilename();
	}
	else
	{
	imageBase = selectedProcess->imageBase;
	fileName = selectedProcess->filename;
	}

	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("\tImagebase: ")TEXT(PRINTF_DWORD_PTR_FULL), imageBase);
	StatusBar.SetText(PART_IMAGEBASE, stringBuffer);
	StatusBar.SetText(PART_MODULE, fileName);
	StatusBar.SetTipText(PART_MODULE, fileName);
	}
	else
	{
	StatusBar.SetText(PART_IMAGEBASE, L"");
	StatusBar.SetText(PART_MODULE, L"");
	}
	}

	bool MainGui::showFileDialog(WCHAR * selectedFile, bool save, const WCHAR * defFileName, const WCHAR * filter, const WCHAR * defExtension, const WCHAR * directory)
	{
	OPENFILENAME ofn = {0};

	// WTL doesn't support new explorer styles on Vista and up
	// This is because it uses a custom hook, we could remove it or derive
	// from CFileDialog but this solution is easier and allows more control anyway (e.g. initial dir)

	if(defFileName)
	{
	wcscpy_s(selectedFile, MAX_PATH, defFileName);
	}
	else
	{
	selectedFile[0] = _T('\0');
	}

	ofn.lStructSize = sizeof(ofn);
	ofn.hwndOwner = m_hWnd;
	ofn.lpstrFilter = filter;
	ofn.lpstrDefExt = defExtension; // only first 3 chars are used, no dots!
	ofn.lpstrFile = selectedFile;
	ofn.lpstrInitialDir = directory;
	ofn.nMaxFile = MAX_PATH;
	ofn.Flags = OFN_PATHMUSTEXIST \| OFN_HIDEREADONLY;

	/*
	*OFN_EXPLORER is automatically used, it only has to be specified
	*if using a custom hook
	*OFN_LONGNAMES is automatically used by explorer-style dialogs
	*/

	if(save)
	ofn.Flags \|= OFN_OVERWRITEPROMPT;
	else
	ofn.Flags \|= OFN_FILEMUSTEXIST;

	if(save)
	return 0 != GetSaveFileName(&ofn);
	else
	return 0 != GetOpenFileName(&ofn);
	}

	void MainGui::setIconAndDialogCaption()
	{
	SetIcon(hIcon, TRUE);
	SetIcon(hIcon, FALSE);

	SetWindowText(TEXT(APPNAME)TEXT(" ")TEXT(ARCHITECTURE)TEXT(" ")TEXT(APPVERSION));
	}

	void MainGui::pickDllActionHandler()
	{
	if(!selectedProcess)
	return;

	PickDllGui dlgPickDll(processAccessHelp.moduleList);
	if(dlgPickDll.DoModal())
	{
	//get selected module
	processAccessHelp.selectedModule = dlgPickDll.getSelectedModule();
	+
	+ processAccessHelp.targetImageBase = processAccessHelp.selectedModule->modBaseAddr;
	+
	Logger::printfDialog(TEXT("->>> Module %s selected."), processAccessHelp.selectedModule->getFilename());
	Logger::printfDialog(TEXT("Imagebase: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: %08X"),processAccessHelp.selectedModule->modBaseAddr,processAccessHelp.selectedModule->modBaseSize);
	}
	else
	{
	processAccessHelp.selectedModule = 0;
	}

	fillStatusBar();
	}

	void MainGui::pickApiActionHandler(CTreeItem item)
	{
	if(!importsHandling.isImport(item))
	return;

	// TODO: new node when user picked an API from another DLL?

	PickApiGui dlgPickApi(processAccessHelp.moduleList);
	if(dlgPickApi.DoModal())
	{
	const ApiInfo* api = dlgPickApi.getSelectedApi();
	if(api && api->module)
	{
	importsHandling.setImport(item, api->module->getFilename(), api->name, api->ordinal, api->hint, true, api->isForwarded);
	}
	}

	fillStatusBar();
	}

	void MainGui::startDisassemblerGui(CTreeItem selectedTreeNode)
	{
	if(!selectedProcess)
	return;

	DWORD_PTR address = importsHandling.getApiAddressByNode(selectedTreeNode);
	if (address)
	{
	BYTE test;
	if(!ProcessAccessHelp::readMemoryFromProcess(address, sizeof(test), &test))
	{
	swprintf_s(stringBuffer, _countof(stringBuffer), TEXT("Can't read memory at ")TEXT(PRINTF_DWORD_PTR_FULL),address);
	MessageBox(stringBuffer, L"Failure", MB_ICONERROR);
	}
	else
	{
	DisassemblerGui dlgDisassembler(address);
	dlgDisassembler.DoModal();
	}
	}
	}

	void MainGui::processSelectedActionHandler(int index)
	{
	std::vector<Process>& processList = processLister.getProcessList();
	Process &process = processList.at(index);
	selectedProcess = 0;

	clearImportsActionHandler();

	Logger::printfDialog(TEXT("Analyzing %s"),process.fullPath);

	if (processAccessHelp.hProcess != 0)
	{
	processAccessHelp.closeProcessHandle();
	apiReader.clearAll();
	}

	if (!processAccessHelp.openProcessHandle(process.PID))
	{
	enableDialogControls(FALSE);
	Logger::printfDialog(TEXT("Error: Cannot open process handle."));
	fillStatusBar();
	return;
	}

	processAccessHelp.getProcessModules(process.PID, processAccessHelp.moduleList);

	apiReader.readApisFromModuleList();

	Logger::printfDialog(TEXT("Loading modules done."));

	//TODO improve
	processAccessHelp.selectedModule = 0;
	processAccessHelp.targetSizeOfImage = process.imageSize;
	processAccessHelp.targetImageBase = process.imageBase;

	ProcessAccessHelp::getSizeOfImageCurrentProcess();

	process.imageSize = (DWORD)processAccessHelp.targetSizeOfImage;


	Logger::printfDialog(TEXT("Imagebase: ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size: %08X"),process.imageBase, process.imageSize);

	process.entryPoint = ProcessAccessHelp::getEntryPointFromFile(process.fullPath);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),process.entryPoint + process.imageBase);

	EditOEPAddress.SetWindowText(stringBuffer);

	selectedProcess = &process;
	enableDialogControls(TRUE);

	fillStatusBar();
	}

	void MainGui::fillProcessListComboBox(CComboBox& hCombo)
	{
	hCombo.ResetContent();

	std::vector<Process>& processList = processLister.getProcessListSnapshot();

	for (size_t i = 0; i < processList.size(); i++)
	{
	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("0x%04X - %s - %s"),processList[i].PID,processList[i].filename,processList[i].fullPath);
	hCombo.AddString(stringBuffer);
	}
	}

	void MainGui::addTextToOutputLog(const WCHAR * text)
	{
	if (m_hWnd)
	{
	ListLog.SetCurSel(ListLog.AddString(text));
	}
	}

	void MainGui::clearOutputLog()
	{
	if (m_hWnd)
	{
	ListLog.ResetContent();
	}
	}

	bool MainGui::saveLogToFile(const WCHAR * file)
	{
	const BYTE BOM[] = {0xFF, 0xFE}; // UTF-16 little-endian
	const WCHAR newLine[] = L"\r\n";
	bool success = true;

	HANDLE hFile = CreateFile(file, GENERIC_WRITE, FILE_SHARE_READ, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
	if(hFile != INVALID_HANDLE_VALUE)
	{
	ProcessAccessHelp::writeMemoryToFileEnd(hFile, sizeof(BOM), BOM);

	WCHAR * buffer = 0;
	int bufsize = 0;
	for(int i = 0; i < ListLog.GetCount(); i++)
	{
	int size = ListLog.GetTextLen(i);
	size += _countof(newLine)-1;
	if(size+1 > bufsize)
	{
	bufsize = size+1;
	delete[] buffer;
	try
	{
	buffer = new WCHAR[bufsize];
	}
	catch(std::bad_alloc&)
	{
	buffer = 0;
	success = false;
	break;
	}
	}

	ListLog.GetText(i, buffer);
	wcscat_s(buffer, bufsize, newLine);

	ProcessAccessHelp::writeMemoryToFileEnd(hFile, size * sizeof(WCHAR), buffer);
	}
	delete[] buffer;
	CloseHandle(hFile);
	}
	return success;
	}

	void MainGui::showInvalidImportsActionHandler()
	{
	importsHandling.selectImports(true, false);
	}

	void MainGui::showSuspectImportsActionHandler()
	{
	importsHandling.selectImports(false, true);
	}

	void MainGui::iatAutosearchActionHandler()
	{
	DWORD_PTR searchAddress = 0;
	DWORD_PTR addressIAT = 0;
	DWORD sizeIAT = 0;
	IATSearch iatSearch;

	if(!selectedProcess)
	return;

	if(EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 1)
	{
	searchAddress = stringToDwordPtr(stringBuffer);
	if (searchAddress)
	{
	if (iatSearch.searchImportAddressTableInProcess(searchAddress, &addressIAT, &sizeIAT))
	{
	Logger::printfDialog(TEXT("IAT found at VA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" RVA ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d)"),addressIAT, addressIAT - processAccessHelp.targetImageBase,sizeIAT,sizeIAT);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT(PRINTF_DWORD_PTR_FULL),addressIAT);
	EditIATAddress.SetWindowText(stringBuffer);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("%08X"),sizeIAT);
	EditIATSize.SetWindowText(stringBuffer);

	swprintf_s(stringBuffer, _countof(stringBuffer),TEXT("IAT found! Start Address ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT(" Size 0x%04X (%d) "),addressIAT,sizeIAT,sizeIAT);
	MessageBox(stringBuffer, L"IAT found", MB_ICONINFORMATION);
	}
	else
	{
	Logger::printfDialog(TEXT("IAT not found at OEP ")TEXT(PRINTF_DWORD_PTR_FULL)TEXT("!"),searchAddress);
	}
	}
	}
	}

	void MainGui::getImportsActionHandler()
	{
	DWORD_PTR addressIAT = 0;
	DWORD sizeIAT = 0;

	if(!selectedProcess)
	return;

	if (EditIATAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
	{
	addressIAT = stringToDwordPtr(stringBuffer);
	}

	if (EditIATSize.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
	{
	sizeIAT = wcstoul(stringBuffer, NULL, 16);
	}

	if (addressIAT && sizeIAT)
	{
	apiReader.readAndParseIAT(addressIAT, sizeIAT,importsHandling.moduleList);
	importsHandling.displayAllImports();
	}

	fillStatusBar();
	}

	DWORD_PTR MainGui::stringToDwordPtr(const WCHAR * hexString)
	{
	DWORD_PTR address = 0;

	#ifdef _WIN64
	address = _wcstoui64(hexString, NULL, 16);
	#else
	address = wcstoul(hexString, NULL, 16);
	#endif

	if (address == 0)
	{
	#ifdef DEBUG_COMMENTS
	Logger::debugLog(L"stringToDwordPtr :: address == 0, %s",hexString);
	#endif
	return 0;
	}
	else
	{
	return address;
	}
	}

	void MainGui::SetupImportsMenuItems(bool isItem, bool isThunk)
	{
	// assert(!(!isItem && isThunk));

	CMenuHandle hSub = hMenuImports.GetSubMenu(0);

	UINT itemOnly = isItem ? MF_ENABLED : MF_GRAYED;
	UINT thunkOnly = isThunk ? MF_ENABLED : MF_GRAYED;

	hSub.EnableMenuItem(ID__INVALIDATEFUNCTION, thunkOnly);
	hSub.EnableMenuItem(ID__DISASSEMBLE, thunkOnly);
	hSub.EnableMenuItem(ID__CUTTHUNK, thunkOnly);

	hSub.EnableMenuItem(ID__DELETETREENODE, itemOnly);
	}

	void MainGui::DisplayContextMenuImports(CWindow hwnd, CPoint pt)
	{
	if(TreeImports.GetCount() < 1)
	return;

	CTreeItem over, parent;

	if(pt.x == -1 && pt.y == -1) // invoked by keyboard
	{
	CRect pos;
	over = TreeImports.GetFirstSelectedItem();
	if(over)
	{
	over.EnsureVisible();
	over.GetRect(&pos, TRUE);
	TreeImports.ClientToScreen(&pos);
	}
	else
	{
	TreeImports.GetWindowRect(&pos);
	}
	pt = pos.TopLeft();
	}
	else
	{
	// Get item under cursor
	CPoint client = pt;
	TreeImports.ScreenToClient(&client);
	UINT flags;
	over = TreeImports.HitTest(client, &flags);
	if(over && !(flags & TVHT_ONITEM))
	{
	over = NULL;
	}
	}

	if (hMenuImports)
	{
	// Prepare hmenuImports
	SetupImportsMenuItems(!over.IsNull(), importsHandling.isImport(over));

	CMenuHandle hSub = hMenuImports.GetSubMenu(0);

	BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN \| TPM_RIGHTBUTTON \| TPM_RETURNCMD, pt.x, pt.y, hwnd);
	if (menuItem)
	{
	if ((menuItem >= PLUGIN_MENU_BASE_ID) && (menuItem <= (int)(PluginLoader::getScyllaPluginList().size() + PluginLoader::getImprecPluginList().size() + PLUGIN_MENU_BASE_ID)))
	{
	//wsprintf(stringBuffer, L"%d %s\n",menuItem,pluginList[menuItem - PLUGIN_MENU_BASE_ID].pluginName);
	//MessageBox(stringBuffer, L"plugin selection");

	pluginActionHandler(menuItem);
	return;
	}
	switch (menuItem)
	{
	case ID__INVALIDATEFUNCTION:
	importsHandling.invalidateImport(over);
	break;
	case ID__DISASSEMBLE:
	startDisassemblerGui(over);
	break;
	case ID__EXPANDALLNODES:
	importsHandling.expandAllTreeNodes();
	break;
	case ID__COLLAPSEALLNODES:
	importsHandling.collapseAllTreeNodes();
	break;
	case ID__CUTTHUNK:
	importsHandling.cutImport(over);
	break;
	case ID__DELETETREENODE:
	importsHandling.cutModule(importsHandling.isImport(over) ? over.GetParent() : over);
	break;
	}
	}
	}

	fillStatusBar(); // ?
	}

	void MainGui::DisplayContextMenuLog(CWindow hwnd, CPoint pt)
	{
	if (hMenuLog)
	{
	if(pt.x == -1 && pt.y == -1) // invoked by keyboard
	{
	CRect pos;
	ListLog.GetWindowRect(&pos);
	pt = pos.TopLeft();
	}

	CMenuHandle hSub = hMenuLog.GetSubMenu(0);
	BOOL menuItem = hSub.TrackPopupMenu(TPM_LEFTALIGN \| TPM_RIGHTBUTTON \| TPM_RETURNCMD, pt.x, pt.y, hwnd);
	if (menuItem)
	{
	switch (menuItem)
	{
	case ID__SAVE:
	WCHAR selectedFilePath[MAX_PATH];
	if(showFileDialog(selectedFilePath, true, NULL, filterTxt, L"txt"))
	{
	saveLogToFile(selectedFilePath);
	}
	break;
	case ID__CLEAR:
	clearOutputLog();
	break;
	}
	}
	}
	}

	void MainGui::appendPluginListToMenu(CMenuHandle hMenu)
	{
	std::vector<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList();
	std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList();

	if (scyllaPluginList.size() > 0)
	{
	CMenuHandle newMenu;
	newMenu.CreatePopupMenu();

	for (size_t i = 0; i < scyllaPluginList.size(); i++)
	{
	newMenu.AppendMenu(MF_STRING, i + PLUGIN_MENU_BASE_ID, scyllaPluginList[i].pluginName);
	}

	hMenu.AppendMenu(MF_MENUBARBREAK);
	hMenu.AppendMenu(MF_POPUP, newMenu, L"Scylla Plugins");
	}

	if (imprecPluginList.size() > 0)
	{
	CMenuHandle newMenu;
	newMenu.CreatePopupMenu();

	for (size_t i = 0; i < imprecPluginList.size(); i++)
	{
	newMenu.AppendMenu(MF_STRING, scyllaPluginList.size() + i + PLUGIN_MENU_BASE_ID, imprecPluginList[i].pluginName);
	}

	hMenu.AppendMenu(MF_MENUBARBREAK);
	hMenu.AppendMenu(MF_POPUP, newMenu, L"ImpREC Plugins");
	}
	}

	void MainGui::dumpActionHandler()
	{
	if(!selectedProcess)
	return;

	WCHAR selectedFilePath[MAX_PATH];
	const WCHAR * fileFilter;
	const WCHAR * defExtension;
	PeDump peDump;

	if (processAccessHelp.selectedModule)
	{
	fileFilter = filterDll;
	defExtension = L"dll";
	}
	else
	{
	fileFilter = filterExe;
	defExtension = L"exe";
	}

	if(showFileDialog(selectedFilePath, true, NULL, fileFilter, defExtension))
	{
	if (processAccessHelp.selectedModule)
	{
	//dump DLL

	peDump.imageBase = processAccessHelp.selectedModule->modBaseAddr;
	peDump.sizeOfImage = processAccessHelp.selectedModule->modBaseSize;
	//get it from gui
	peDump.entryPoint = getOEPFromGui();
	wcscpy_s(peDump.fullpath, MAX_PATH, processAccessHelp.selectedModule->fullPath);
	}
	else
	{
	peDump.imageBase = ProcessAccessHelp::targetImageBase;
	peDump.sizeOfImage = (DWORD)ProcessAccessHelp::targetSizeOfImage;
	//get it from gui
	peDump.entryPoint = getOEPFromGui();
	wcscpy_s(peDump.fullpath, MAX_PATH, selectedProcess->fullPath);
	}

	peDump.useHeaderFromDisk = ConfigurationHolder::getConfigObject(USE_PE_HEADER_FROM_DISK)->isTrue();
	if (peDump.dumpCompleteProcessToDisk(selectedFilePath))
	{
	Logger::printfDialog(TEXT("Dump success %s"),selectedFilePath);
	//MessageBox(L"Image dumped successfully.", L"Success");
	}
	else
	{
	Logger::printfDialog(TEXT("Error: Cannot dump image."));
	MessageBox(L"Cannot dump image.", L"Failure", MB_ICONERROR);
	}
	}
	}

	DWORD_PTR MainGui::getOEPFromGui()
	{
	if (EditOEPAddress.GetWindowText(stringBuffer, _countof(stringBuffer)) > 0)
	{
	return stringToDwordPtr(stringBuffer);
	}
	else
	{
	return 0;
	}
	}

	void MainGui::peRebuildActionHandler()
	{
	DWORD newSize = 0;
	WCHAR selectedFilePath[MAX_PATH];
	PeRebuild peRebuild;

	if(showFileDialog(selectedFilePath, false, NULL, filterExeDll))
	{
	if (ConfigurationHolder::getConfigObject(CREATE_BACKUP)->isTrue())
	{
	if (!ProcessAccessHelp::createBackupFile(selectedFilePath))
	{
	Logger::printfDialog(TEXT("Creating backup file failed %s"), selectedFilePath);
	}
	}

	LONGLONG fileSize = ProcessAccessHelp::getFileSize(selectedFilePath);
	LPVOID mapped = peRebuild.createFileMappingViewFull(selectedFilePath);

	newSize = peRebuild.realignPE(mapped, (DWORD)fileSize);
	peRebuild.closeAllMappingHandles();

	if (newSize < 10)
	{
	Logger::printfDialog(TEXT("Rebuild failed %s"), selectedFilePath);
	MessageBox(L"Rebuild failed.", L"Failure", MB_ICONERROR);
	}
	else
	{
	peRebuild.truncateFile(selectedFilePath, newSize);

	Logger::printfDialog(TEXT("Rebuild success %s"), selectedFilePath);
	Logger::printfDialog(TEXT("-> Old file size 0x%08X new file size 0x%08X (%d %%)"), (DWORD)fileSize, newSize, (DWORD)((newSize * 100) / (DWORD)fileSize) );
	//MessageBox(L"Image rebuilded successfully.", L"Success", MB_ICONINFORMATION);
	}
	}
	}

	void MainGui::dumpFixActionHandler()
	{
	if(!selectedProcess)
	return;

	if (TreeImports.GetCount() < 2)
	{
	Logger::printfDialog(TEXT("Nothing to rebuild"));
	return;
	}

	WCHAR newFilePath[MAX_PATH];
	WCHAR selectedFilePath[MAX_PATH];
	const WCHAR * fileFilter;

	ImportRebuild importRebuild;

	if (processAccessHelp.selectedModule)
	{
	fileFilter = filterDll;
	}
	else
	{
	fileFilter = filterExe;
	}

	if (showFileDialog(selectedFilePath, false, NULL, fileFilter))
	{
	wcscpy_s(newFilePath,MAX_PATH,selectedFilePath);

	const WCHAR * extension = 0;

	WCHAR* dot = wcsrchr(newFilePath, L'.');
	if (dot)
	{
	*dot = L'\0';
	extension = selectedFilePath + (dot - newFilePath); //wcsrchr(selectedFilePath, L'.');
	}

	wcscat_s(newFilePath, MAX_PATH, L"_SCY");

	if(extension)
	{
	wcscat_s(newFilePath, MAX_PATH, extension);
	}

	if (importRebuild.rebuildImportTable(selectedFilePath,newFilePath,importsHandling.moduleList))
	{
	//MessageBox(L"Imports rebuilding successful", L"Success", MB_ICONINFORMATION);

	Logger::printfDialog(TEXT("Import Rebuild success %s"), newFilePath);
	}
	else
	{
	Logger::printfDialog(TEXT("Import Rebuild failed, target %s"), selectedFilePath);
	MessageBox(L"Imports rebuilding failed", L"Failure", MB_ICONERROR);
	}
	}
	}

	void MainGui::enableDialogControls(BOOL value)
	{
	BOOL valButton = value ? TRUE : FALSE;

	GetDlgItem(IDC_BTN_PICKDLL).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_DUMP).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_FIXDUMP).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_IATAUTOSEARCH).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_GETIMPORTS).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_SUSPECTIMPORTS).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_INVALIDIMPORTS).EnableWindow(valButton);
	GetDlgItem(IDC_BTN_CLEARIMPORTS).EnableWindow(valButton);

	CMenuHandle menu = GetMenu();

	UINT valMenu = value ? MF_ENABLED : MF_GRAYED;

	menu.EnableMenuItem(ID_FILE_DUMP, valMenu);
	menu.EnableMenuItem(ID_FILE_FIXDUMP, valMenu);
	menu.EnableMenuItem(ID_IMPORTS_INVALIDATESELECTED, valMenu);
	menu.EnableMenuItem(ID_IMPORTS_CUTSELECTED, valMenu);
	menu.EnableMenuItem(ID_MISC_DLLINJECTION, valMenu);
	menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetScylla, MF_BYPOSITION \| valMenu);
	menu.GetSubMenu(MenuImportsOffsetTrace).EnableMenuItem(MenuImportsTraceOffsetImpRec, MF_BYPOSITION \| valMenu);

	//not yet implemented
	GetDlgItem(IDC_BTN_SAVETREE).EnableWindow(FALSE);
	GetDlgItem(IDC_BTN_LOADTREE).EnableWindow(FALSE);
	GetDlgItem(IDC_BTN_AUTOTRACE).EnableWindow(FALSE);

	menu.EnableMenuItem(ID_IMPORTS_SAVETREE, MF_GRAYED);
	menu.EnableMenuItem(ID_IMPORTS_LOADTREE, MF_GRAYED);
	menu.EnableMenuItem(ID_TRACE_AUTOTRACE, MF_GRAYED);
	}

	void MainGui::showAboutDialog()
	{
	AboutGui dlgAbout;
	dlgAbout.DoModal();
	}

	void MainGui::dllInjectActionHandler()
	{
	if(!selectedProcess)
	return;

	WCHAR selectedFilePath[MAX_PATH];
	HMODULE hMod = 0;
	DllInjection dllInjection;

	if (showFileDialog(selectedFilePath, false, NULL, filterDll))
	{
	hMod = dllInjection.dllInjection(ProcessAccessHelp::hProcess, selectedFilePath);
	if (hMod && ConfigurationHolder::getConfigObject(DLL_INJECTION_AUTO_UNLOAD)->isTrue())
	{
	if (!dllInjection.unloadDllInProcess(ProcessAccessHelp::hProcess, hMod))
	{
	Logger::printfDialog(TEXT("DLL unloading failed, target %s"), selectedFilePath);
	}
	}

	if (hMod)
	{
	Logger::printfDialog(TEXT("DLL Injection was successful, target %s"), selectedFilePath);
	}
	else
	{
	Logger::printfDialog(TEXT("DLL Injection failed, target %s"), selectedFilePath);
	}
	}
	}

	void MainGui::optionsActionHandler()
	{
	OptionsGui dlgOptions;
	dlgOptions.DoModal();
	}

	void MainGui::clearImportsActionHandler()
	{
	importsHandling.clearAllImports();
	fillStatusBar();
	}

	void MainGui::pluginActionHandler( int menuItem )
	{
	if(!selectedProcess)
	return;

	DllInjectionPlugin dllInjectionPlugin;

	std::vector<Plugin> &scyllaPluginList = PluginLoader::getScyllaPluginList();
	std::vector<Plugin> &imprecPluginList = PluginLoader::getImprecPluginList();

	menuItem -= PLUGIN_MENU_BASE_ID;

	dllInjectionPlugin.hProcess = ProcessAccessHelp::hProcess;
	dllInjectionPlugin.apiReader = &apiReader;

	if (menuItem < (int)scyllaPluginList.size())
	{
	//scylla plugin
	dllInjectionPlugin.injectPlugin(scyllaPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize);
	}
	else
	{
	#ifndef _WIN64

	menuItem -= (int)scyllaPluginList.size();
	//imprec plugin
	dllInjectionPlugin.injectImprecPlugin(imprecPluginList[menuItem], importsHandling.moduleList,selectedProcess->imageBase, selectedProcess->imageSize);

	#endif
	}

	importsHandling.scanAndFixModuleList();
	importsHandling.displayAllImports();
	fillStatusBar();
	}

File Metadata

Mime Type: text/x-diff
Expires: Tue, Apr 15, 7:07 PM (5 h, 26 m)
Storage Engine: local-disk
Storage Format: Raw Data
Storage Handle: f9/58/171bffce760eb5bd4a88d374ca5c

No OneTemporaryActions

View Options

File Metadata

Event Timeline

No OneTemporary
Actions