#include "StdAfx.h"
#define DONT_INCLUDE_AFXTEMPL
#include "CPEFile.h"

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

CPEFile::CPEFile() {
	mFile = NULL;
}

CPEFile::~CPEFile() {
	Close();
}

/*virtual*/ void CPEFile::Close() {
	if (mFile) {
		mFile->Close();
		delete mFile;
	}
	mFile = NULL;

	if (mSectionHeaders)
		free(mSectionHeaders);
	mSectionHeaders = NULL;
}

/*virtual*/ BOOL CPEFile::Open(const char*filename, unsigned int flags, CFileException*pError) {
	BOOL result = false;

#line 36
	mFile = new CFile;
	if (mFile->Open(filename, flags, pError)) {
		result = Init();
		if (!result)
			Close();
	}

	return result;
}

BOOL CPEFile::Init() {
	IMAGE_DOS_HEADER header;

	mSectionHeaders = NULL;

	memset(&header, 0, sizeof(header));
	mFile->Read(&header, sizeof(header));
	if (header.e_magic != IMAGE_DOS_SIGNATURE)
		return false;

	mFile->Seek(header.e_lfanew, CFile::begin);

	DWORD magic = 0;
	mFile->Read(&magic, sizeof(magic));
	if (magic != IMAGE_NT_SIGNATURE)
		return false;

	memset(&mFileHeader, 0, sizeof(mFileHeader));
	mFile->Read(&mFileHeader, sizeof(mFileHeader));
	mFile->Seek(IMAGE_SIZEOF_NT_OPTIONAL_HEADER, CFile::current);

	mSectionHeaders = (IMAGE_SECTION_HEADER *) malloc(sizeof(IMAGE_SECTION_HEADER) * mFileHeader.NumberOfSections);
	mFile->Read(mSectionHeaders, sizeof(IMAGE_SECTION_HEADER) * mFileHeader.NumberOfSections);

	int i;
	for (i = 0; i < mFileHeader.NumberOfSections; i++) {
		if (strcmp((const char *) mSectionHeaders[i].Name, ".rsrc") == 0)
			break;
	}

	if (i == mFileHeader.NumberOfSections)
		return false;

	mRsrcPtr = mSectionHeaders[i].PointerToRawData;
	mRsrcIndex = i;
	return true;
}

BOOL CPEFile::EnumResource(CString& s, RESTRACKERWORK& work) {
	IMAGE_RESOURCE_DIRECTORY_ENTRY entry;

	if (!EnumResourceDirectory(&entry, work))
		return false;

	GetResourceName(s, entry);
	return true;
}

BOOL CPEFile::EnumResourceDirectory(IMAGE_RESOURCE_DIRECTORY_ENTRY* entry, RESTRACKERWORK& work) {
	if (!work.didRead) {
		mFile->Seek(mRsrcPtr, CFile::begin);
		mFile->Read(&work.directory, sizeof(work.directory));
		work.position = mFile->GetPosition();
		work.didRead = true;
		work.nextIndex = 0;
	}

	int entryCount = (work.directory.NumberOfNamedEntries + work.directory.NumberOfIdEntries);
	if (work.nextIndex >= entryCount)
		return false;

	mFile->Seek(work.position + work.nextIndex * sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY), CFile::begin);
	mFile->Read(entry, sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY));
	memcpy(&work.entry, entry, sizeof(IMAGE_RESOURCE_DIRECTORY_ENTRY));
	work.nextIndex++;
	return true;
}

void CPEFile::GetResourceName(CString& s, IMAGE_RESOURCE_DIRECTORY_ENTRY& entry) {
	if (entry.Name & 0x80000000) {
		unsigned short nameLen;
		mFile->Seek(mRsrcPtr + entry.NameOffset, CFile::begin);
		mFile->Read(&nameLen, sizeof(nameLen));

		char *wstr = (char *) malloc(nameLen * 2);
		mFile->Read(wstr, nameLen * 2);

		int bufferSize = nameLen * 2 + 1;
		char *theBuf = (char *) malloc(bufferSize);
		memset(theBuf, 0, bufferSize);

		int i;
		for (i = 0; i < nameLen; i++) {
			theBuf[i] = wstr[i * 2];
		}
		theBuf[i] = 0;

		s = theBuf;
		free(wstr);
		free(theBuf);
	} else {
		s.Format("#%u", entry.Name);
	}

	if (entry.OffsetToData & 0x80000000)
		s += "\\";
}

void CPEFile::EnterDirectory(RESTRACKERWORK& child, RESTRACKERWORK& parent) {
	mFile->Seek(mRsrcPtr + parent.entry.OffsetToDirectory, CFile::begin);
	mFile->Read(&child.directory, sizeof(child.directory));
	child.position = mFile->GetPosition();
	child.didRead = true;
	child.nextIndex = 0;
}

/*virtual*/ BOOL CPEFile::Seek(const char* inName) {
#line 141
	_ASSERT(strncmp(inName, "\\.rsrc", 6) == 0);

	char *name = (char *) malloc(strlen(inName) + 1);
	strcpy(name, inName);

	CString elem;
	elem = (char *) NULL;
	elem = strtok(name, "\\");

	RESTRACKERWORK work;
	memset(&work, 0, sizeof(work));

	BOOL isOK = false;

	CString resName = "";

	while ((elem = strtok(NULL, "\\")) != "") {
		if (isOK) {
			RESTRACKERWORK child;
			EnterDirectory(child, work);
			work = child;
		}

		isOK = false;
		while (EnumResource(resName, work) == 1) {
			if (_stricmp(resName, elem) == 0 || _stricmp(resName, elem + "\\") == 0) {
				isOK = true;
				break;
			}
		}

		if (!isOK)
			return false;
	}

	free(name);

	if (resName.Right(1) == "\\")
		return false;

	IMAGE_RESOURCE_DATA_ENTRY theEntryData;

	mFile->Seek(mRsrcPtr + work.entry.OffsetToData, CFile::begin);
	mFile->Read(&theEntryData, sizeof(theEntryData));

	mStartPosition = theEntryData.OffsetToData - (mSectionHeaders[mRsrcIndex].VirtualAddress - mSectionHeaders[mRsrcIndex].PointerToRawData);
	mLength = theEntryData.Size;

	mFile->Seek(mStartPosition, CFile::begin);

	return true;
}

/*virtual*/ unsigned long CPEFile::GetLength() {
	return mLength;
}

/*virtual*/ unsigned long CPEFile::GetPosition() {
	return mFile->GetPosition() - mStartPosition;
}

/*virtual*/ long CPEFile::Seek(long offset, unsigned int from) {
	if (from == CFile::begin)
		return mFile->Seek(offset + mStartPosition, from) - mStartPosition;
	else
		return mFile->Seek(offset, from) - mStartPosition;
}

/*virtual*/ unsigned int CPEFile::Read(void* buf, unsigned int len) {
	int remaining = mLength - GetPosition();
	if (remaining < 0 || len > remaining)
		len = remaining;
	return mFile->Read(buf, len);
}

/*virtual*/ void CPEFile::Write(void* buf, unsigned int len) {
	int remaining = mLength - GetPosition();
	if (remaining < 0 || len > remaining)
		MessageBox(NULL, "CPEFile::Write overrun", "ERROR", MB_OK);
	else
		mFile->Write(buf, len);
}