#include <common/pathfnd.hpp>
#include <common/pvector.hpp>
#include <common/pvector.tpp>
#include <common/pview.hpp>
#include <common/openfile.hpp>
#include <image/pehdr.hpp>
#include <image/dbgdir.hpp>

WORD PEHeader::operator<<(PureViewOfFile &pureView)
{
	pureView.rewind();
	if(!((DOSHeader&)*this<<pureView))return FALSE;
	if(!pureView.seek(neHeaderAddress(),PureViewOfFile::SeekSet))return FALSE;
	pureView.read(mPESignature);
	if(!((ImageHeader&)*this<<pureView))return FALSE;
	if(sizeof(ImageOptionalHeader)==sizeOptionalHeader())
		if(!((ImageOptionalHeader&)*this<<pureView))return FALSE;
	size(sectionCount());
	for(short itemIndex=0;itemIndex<sectionCount();itemIndex++)
		((PureVector<ImageSectionHeader>&)*this)[itemIndex]<<pureView;
	return TRUE;
}

WORD PEHeader::loadImageImportDescriptors(Block<ImageImportDescriptor> &imageImportDescriptors,ImageSectionHeader &importSectionHeader,PureViewOfFile &pureView,WORD memImage)
{
	DirectoryEntryKey sectionKey(DirectoryEntryKey::ImportDirectory);
	ImageDataDirectory imageSectionDirectory(((ImageOptionalHeader&)*this)[sectionKey]);
	ImageImportDescriptor imageImportDescriptor;
	ImageSectionHeader importNamesSectionHeader;
	DWORD deltaOffset;
	DWORD deltaNames;

	imageImportDescriptors.remove();
  isMemImage(memImage);
	if(!enclosingSectionHeader(imageSectionDirectory.virtualAddress(),importSectionHeader,deltaOffset))return FALSE;
	pureView.seek(imageSectionDirectory.virtualAddress()-deltaOffset,PureViewOfFile::SeekSet);
	while(TRUE)
	{
		String moduleName;
		imageImportDescriptor<<pureView;
		if(!imageImportDescriptor.isOkay())break;
		enclosingSectionHeader(imageImportDescriptor.nameRVA(),importNamesSectionHeader,deltaNames);
		pureView.push();
		pureView.seek(imageImportDescriptor.nameRVA()-deltaNames,PureViewOfFile::SeekSet);
		pureView.getLine(moduleName);
		imageImportDescriptor.moduleName(moduleName);
		pureView.pop();
		imageImportDescriptor.loadImageImportThunks(deltaOffset,pureView,(ImageSectionHeaders&)*this);
		imageImportDescriptors.insert(&imageImportDescriptor);
	}
	return imageImportDescriptors.size();
}

WORD PEHeader::loadImageFirstThunk(Block<ImageImportDescriptor> &imageImportDescriptors,ImageSectionHeader &importSectionHeader,PureViewOfFile &pureView,WORD memImage)
{
	DirectoryEntryKey sectionKey(DirectoryEntryKey::ImportDirectory);
	ImageDataDirectory imageSectionDirectory(((ImageOptionalHeader&)*this)[sectionKey]);
	ImageImportDescriptor imageImportDescriptor;
	ImageSectionHeader importNamesSectionHeader;
	DWORD deltaOffset;
	DWORD deltaNames;

	imageImportDescriptors.remove();
  isMemImage(memImage);
	if(!enclosingSectionHeader(imageSectionDirectory.virtualAddress(),importSectionHeader,deltaOffset))return FALSE;
	pureView.seek(imageSectionDirectory.virtualAddress()-deltaOffset,PureViewOfFile::SeekSet);
	while(TRUE)
	{
		String moduleName;
		imageImportDescriptor<<pureView;
		if(!imageImportDescriptor.isOkay())break;
		enclosingSectionHeader(imageImportDescriptor.nameRVA(),importNamesSectionHeader,deltaNames);
		pureView.push();
		pureView.seek(imageImportDescriptor.nameRVA()-deltaNames,PureViewOfFile::SeekSet);
		pureView.getLine(moduleName);
		imageImportDescriptor.moduleName(moduleName);
		pureView.pop();
		imageImportDescriptor.loadImageFirstThunk(deltaOffset,pureView,(ImageSectionHeaders&)*this);
		imageImportDescriptors.insert(&imageImportDescriptor);
	}
	return imageImportDescriptors.size();
}

WORD PEHeader::loadImageExportDescriptors(ImageExportDescriptors &imageExportDescriptors,ImageSectionHeader &exportSectionHeader,PureViewOfFile &pureView,WORD sortFlag,WORD memImage)
{
	DirectoryEntryKey sectionKey(DirectoryEntryKey::ExportDirectory);
	ImageDataDirectory imageSectionDirectory(((ImageOptionalHeader&)*this)[sectionKey]);
	DWORD deltaOffset;

	imageExportDescriptors.size(0);
	if(!imageSectionDirectory.size())return FALSE;
  isMemImage(memImage);
	if(!enclosingSectionHeader(imageSectionDirectory.virtualAddress(),exportSectionHeader,deltaOffset))return FALSE;
	if(mImageExportDescriptorsCache.size())
	{
		imageExportDescriptors=mImageExportDescriptorsCache;
		return (imageExportDescriptors.size()?TRUE:FALSE);
	}
	pureView.push();
	pureView.seek(imageSectionDirectory.virtualAddress()-deltaOffset,PureViewOfFile::SeekSet);
	mImageExportDescriptorsCache.loadImageExports(pureView,(ImageSectionHeaders&)*this,sortFlag);
	imageExportDescriptors=mImageExportDescriptorsCache;
	pureView.pop();
	return (imageExportDescriptors.size()?TRUE:FALSE);
}

WORD PEHeader::loadImageExportDescriptors(ImageExportDescriptors &imageExportDescriptors,ImageSectionHeader &exportSectionHeader,const String &moduleName)
{
	String pathFileName;
	PathFind pathFind;
	PEHeader peHeader;

	if(moduleName.isNull())return FALSE;
	if(!pathFind.findFile(moduleName,pathFileName))return FALSE;
	FileHandle peFile(pathFileName,FileHandle::Read,FileHandle::ShareReadWrite);
	if(!peFile.isOkay())return FALSE;
	FileMap peMap(peFile);
	PureViewOfFile peView(peMap);
	peHeader<<peView;
	peHeader.loadImageExportDescriptors(imageExportDescriptors,exportSectionHeader,peView);
	return imageExportDescriptors.size();
}

WORD PEHeader::loadImageDebugDirectory(Block<ImageDebugDirectory> &imageDebugDirectories,PureViewOfFile &pureView,WORD memImage)
{
	DirectoryEntryKey sectionKey(DirectoryEntryKey::DebugDirectory);
	ImageDataDirectory imageSectionDirectory(((ImageOptionalHeader&)*this)[sectionKey]);
	ImageSectionHeader exportSectionHeader;
	DWORD imageDirectoryEntries;
	DWORD deltaOffset;
	WORD msImage(FALSE);

  imageDebugDirectories.remove();
  isMemImage(memImage);
	if(!enclosingSectionHeader(imageSectionDirectory.virtualAddress(),exportSectionHeader,deltaOffset))return FALSE;
	for(short sectionIndex=0;sectionIndex<size();sectionIndex++)if(((PureVector<ImageSectionHeader>&)*this)[sectionIndex].name()==String(".text"))msImage=TRUE;
	imageDirectoryEntries=imageSectionDirectory.size();
  if(msImage)imageDirectoryEntries/=sizeof(IMAGE_DEBUG_DIRECTORY);
	pureView.push();
	pureView.seek(imageSectionDirectory.virtualAddress()-deltaOffset,PureViewOfFile::SeekSet);
	for(long entryIndex=0;entryIndex<imageDirectoryEntries;entryIndex++)
	{
		imageDebugDirectories.insert(&ImageDebugDirectory());
		imageDebugDirectories[entryIndex]<<pureView;
	}
	pureView.pop();
	return FALSE;
}