TitanEngine/scylla_wrapper/ImportRebuilder.h

101 lines
3.5 KiB
C++

#pragma once
#include <map>
#include "PeParser.h"
#include "Thunks.h"
#include "IATReferenceScan.h"
class ImportRebuilder : public PeParser
{
public:
ImportRebuilder(const WCHAR* file, const WCHAR* sectionName) : PeParser(file, true)
{
pImportDescriptor = 0;
pThunkData = 0;
pImportByName = 0;
numberOfImportDescriptors = 0;
sizeOfImportSection = 0;
sizeOfApiAndModuleNames = 0;
importSectionIndex = 0;
useOFT = false;
sizeOfOFTArray = 0;
newIatInSection = false;
BuildDirectImportsJumpTable = false;
sizeOfJumpTable = 0;
this->sectionName = (WCHAR*)sectionName;
}
ImportRebuilder(const DWORD_PTR iatVA, const DWORD_PTR FileMapVA, const HANDLE hFileMap, const WCHAR* sectionName) : PeParser(iatVA, FileMapVA, hFileMap, true)
{
pImportDescriptor = 0;
pThunkData = 0;
pImportByName = 0;
numberOfImportDescriptors = 0;
sizeOfImportSection = 0;
sizeOfApiAndModuleNames = 0;
importSectionIndex = 0;
useOFT = false;
sizeOfOFTArray = 0;
newIatInSection = false;
BuildDirectImportsJumpTable = false;
sizeOfJumpTable = 0;
this->sectionName = (WCHAR*)sectionName;
}
bool rebuildImportTable(const WCHAR* newFilePath, std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
bool rebuildMappedImportTable(DWORD_PTR iatVA, std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
void enableOFTSupport();
void enableNewIatInSection(DWORD_PTR iatAddress, DWORD iatSize);
int getIATSectionSize(std::map<DWORD_PTR, ImportModuleThunk> & moduleList)
{
this->calculateImportSizes(moduleList);
return (int)this->sizeOfImportSection;
} ;
IATReferenceScan* iatReferenceScan;
bool BuildDirectImportsJumpTable;
private:
PIMAGE_IMPORT_DESCRIPTOR pImportDescriptor;
PIMAGE_THUNK_DATA pThunkData;
PIMAGE_IMPORT_BY_NAME pImportByName;
size_t numberOfImportDescriptors;
size_t sizeOfImportSection;
size_t sizeOfApiAndModuleNames;
size_t importSectionIndex;
WCHAR* sectionName;
//OriginalFirstThunk Array in Import Section
size_t sizeOfOFTArray;
bool useOFT;
bool newIatInSection;
DWORD_PTR IatAddress;
DWORD IatSize;
DWORD sizeOfJumpTable;
DWORD directImportsJumpTableRVA;
BYTE* JMPTableMemory;
DWORD newIatBaseAddressRVA;
DWORD fillImportSection(std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
BYTE* getMemoryPointerFromRVA(DWORD_PTR dwRVA);
bool createNewImportSection(std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
bool buildNewImportTable(std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
bool buildNewMappedImportTable(std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
void setFlagToIATSection(DWORD_PTR iatAddress);
size_t addImportToImportTable(ImportThunk* pImport, PIMAGE_THUNK_DATA pThunk, PIMAGE_IMPORT_BY_NAME pImportByName, DWORD sectionOffset);
size_t addImportDescriptor(ImportModuleThunk* pImportModule, DWORD sectionOffset, DWORD sectionOffsetOFTArray);
void calculateImportSizes(std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
void addSpecialImportDescriptor(DWORD_PTR rvaFirstThunk, DWORD sectionOffsetOFTArray);
void patchFileForNewIatLocation();
void changeIatBaseAddress(std::map<DWORD_PTR, ImportModuleThunk> & moduleList);
void patchFileForDirectImportJumpTable();
};