杀毒

开发平台：
Visual C++

PEFile.cpp：源码内容
							// PEFile.cpp: implementation of the CPEFile class.
//
//////////////////////////////////////////////////////////////////////
#include "stdafx.h"
#include <Imagehlp.h>
//#include "PETools.h"
#include "PEFile.h"
#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif
//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////
CPEFile::CPEFile(HWND hwnd/* = NULL*/)
{
	m_hwnd = hwnd;
	strcpy(m_filename, "");
	m_buffer = 0;
	m_size = 0;
	m_dosstub_size = 0;
	m_pe_header = 0;
	m_std_header = 0;
	m_nt_header = 0;
	m_rom_header = 0;
	m_directories = 0;
	m_sections = 0;
}
CPEFile::~CPEFile()
{
	if (m_buffer)
	{
		delete m_buffer;
		m_buffer = 0;
	}
}
BOOL CheckPe(FILE* pFile)
{		
	fseek(pFile,0,SEEK_SET);
	BOOL	bFlags=FALSE;
	WORD	IsMZ;
	DWORD	IsPE,pNT;
	fread(&IsMZ,sizeof(WORD),1,pFile);
	if(IsMZ==0x5A4D)
	{
		fseek(pFile,0x3c,SEEK_SET);
		fread(&pNT,sizeof(DWORD),1,pFile);
		fseek(pFile,pNT,SEEK_SET);
		fread(&IsPE,sizeof(DWORD),1,pFile);
		if(IsPE==0X00004550)
			bFlags=TRUE;
		else
			bFlags=FALSE;
	}
	else
		bFlags=FALSE;
	fseek(pFile,0,SEEK_SET);
	return bFlags;
}
// LoadExecutable
//
// - Load the executable in memory and parse the pe header
//==============================================================================================
bool CPEFile::LoadExecutable(char *filename)
{
	bool result;
	FILE *f;
    //IMAGE_DOS_HEADER     DosHeader;
	// Delete all previous buffer
	if (m_buffer)
	{
		delete m_buffer;
		m_buffer = 0;
	}
	// Load the file into memory
	f = fopen(filename, "rb");
	if (!f)
	{
		//MessageBox(m_hwnd, CString("Can't open for reading the file ") + CString(filename), "Open error", 0);
		return (false);
	}
	//fseek(f,0,SEEK_SET);
	//fread(&DosHeader,sizeof(IMAGE_DOS_HEADER),1,f);
	//if(DosHeader.e_magic!=IMAGE_DOS_SIGNATURE)
	//{
		//return  false;
	//}
	if(!CheckPe(f))
	{
		return  false;
	}
	fseek(f, 0, SEEK_END);
	m_size = ftell(f);
	fseek(f, 0, SEEK_SET);
	m_buffer = new unsigned char[m_size];
	if (fread(m_buffer, sizeof(unsigned char), m_size, f) != (size_t)m_size)
	{
		MessageBox(m_hwnd, CString("Can't read the whole file ") + CString(filename), "Open error", 0);
		fclose(f);
		if (m_buffer)
		{
			delete m_buffer;
			m_buffer = 0;
		}
		return (false);
	}
	fclose (f);
	// Update all variables
	result = UpdatePEVars(false);
	if (result)
	{
		strcpy(m_filename, filename);
	}
	return (result);
}
// UpdatePEVars
//
// - Parse all pe headers into our variables
//==============================================================================================
bool CPEFile::UpdatePEVars(bool fix_sections/* = false*/)
{
	unsigned int	i;
	unsigned short	dos_stub;
	unsigned int	pe_offset, pe_signature;
	unsigned int	current_pos;
	// DEBUG
	fix_sections = true;
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		//MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	// Check DosStub infos
	dos_stub = ((unsigned short*)m_buffer)[0];
	if (dos_stub != 0x5a4d)
	{
		//MessageBox(m_hwnd, "Not a windows executable", "Format error", 0);
		return (false);
	}
	pe_offset = *((unsigned int*)(m_buffer + 0x3c));
	// Check PE Format signature
	pe_signature = *((unsigned int*)(m_buffer + pe_offset));
	if (pe_signature != 0x00004550)
	{
		//MessageBox(m_hwnd, "Not a PE format executable", "Format error", 0);
		//return (false);
	}
	// Get PE Infos
	current_pos = pe_offset + 4;
	m_dosstub_size = current_pos;
	m_pe_header = (PEHeader*)(m_buffer + current_pos);
	current_pos += sizeof(PEHeader);
	m_nt_header = (NTOptionalHeader*)(m_buffer + current_pos);
	current_pos += sizeof(NTOptionalHeader);
	m_directories = (DataDirectory*)(m_buffer + current_pos);
	current_pos += sizeof(DataDirectory) * m_nt_header->numDataDirectories;
	m_sections = (Section*)(m_buffer + current_pos);
	current_pos += sizeof(Section) * m_pe_header->numSections;
	// Do we need to fix all Raw Infos for each section? (Especially for a dumped task)
	if (fix_sections)
	{
		for (i=0; i<m_pe_header->numSections; i++)
		{
			m_sections[i].dataOffset = m_sections[i].RVA;
			m_sections[i].dataAlignSize = m_sections[i].misc.virtualSize;
		}
	}
	return (true);
}
// WriteInfos
//
// - Write pe infos into a txt file
//==============================================================================================
bool CPEFile::WriteInfos(char *filename)
{
	unsigned int i;
	char buffer[9];
	FILE *fd;
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	
	fd = fopen(filename, "w");
	if (!fd)
	{
		MessageBox(m_hwnd, CString("Can't open for writing the file ") + CString(filename), "Open error", 0);
		return (false);
	}
	// pe header
	fprintf(fd, "---==== PE Infos ====---nn");
	fprintf(fd, "tdosStubSize 0x%xn", m_dosstub_size);
	fprintf(fd, "nPEHeadern");
	fprintf(fd, "tcpuType 0x%xn", m_pe_header->cpuType);
	fprintf(fd, "tnumSections 0x%xn", m_pe_header->numSections);
	fprintf(fd, "tdateStamp 0x%xn", m_pe_header->dateStamp);
	fprintf(fd, "tsymbolTable 0x%xn", m_pe_header->symbolTable);
	fprintf(fd, "tnumSymbols 0x%xn", m_pe_header->numSymbols);
	fprintf(fd, "toptionalHeaderSize 0x%xn", m_pe_header->optionalHeaderSize);
	fprintf(fd, "tflags 0x%xn", m_pe_header->flags);
	// pe optional header
	fprintf(fd, "nNTOptionalHeadern");
	fprintf(fd, "tmagic 0x%xn", m_nt_header->magic);
	fprintf(fd, "tlinkerMajor 0x%xn", m_nt_header->linkerMajor);
	fprintf(fd, "tlinkerMinor 0x%xn", m_nt_header->linkerMinor);
	fprintf(fd, "tcodeSize 0x%xn", m_nt_header->codeSize);
	fprintf(fd, "tinitDataSize 0x%xn", m_nt_header->initDataSize);
	fprintf(fd, "tuninitDataSize 0x%xn", m_nt_header->uninitDataSize);
	fprintf(fd, "tentryPoint 0x%xn", m_nt_header->entryPoint);
	fprintf(fd, "tcodeBase 0x%xn", m_nt_header->codeBase);
	fprintf(fd, "tdataBase 0x%xn", m_nt_header->dataBase);
	fprintf(fd, "timageBase 0x%xn", m_nt_header->imageBase);
	fprintf(fd, "tsectionAlign 0x%xn", m_nt_header->sectionAlign);
	fprintf(fd, "tfileAlign 0x%xn", m_nt_header->fileAlign);
	fprintf(fd, "tosMajor 0x%xn", m_nt_header->osMajor);
	fprintf(fd, "tosMinor 0x%xn", m_nt_header->osMinor);
	fprintf(fd, "timageMajor 0x%xn", m_nt_header->imageMajor);
	fprintf(fd, "timageMinor 0x%xn", m_nt_header->imageMinor);
	fprintf(fd, "tsubsystemMajor 0x%xn", m_nt_header->subsystemMajor);
	fprintf(fd, "tsubsystemMinor 0x%xn", m_nt_header->subsystemMinor);
	fprintf(fd, "treserved 0x%xn", m_nt_header->reserved);
	fprintf(fd, "timageSize 0x%xn", m_nt_header->imageSize);
	fprintf(fd, "theadersSize 0x%xn", m_nt_header->headersSize);
	fprintf(fd, "tchecksum 0x%xn", m_nt_header->checksum);
	fprintf(fd, "tsubsystem 0x%xn", m_nt_header->subsystem);
	fprintf(fd, "tdllFlags 0x%xn", m_nt_header->dllFlags);
	fprintf(fd, "tstackReserveSize 0x%xn", m_nt_header->stackReserveSize);
	fprintf(fd, "tstackCommitSize 0x%xn", m_nt_header->stackCommitSize);
	fprintf(fd, "theapReserveSize 0x%xn", m_nt_header->heapReserveSize);
	fprintf(fd, "theapCommitSize 0x%xn", m_nt_header->heapCommitSize);
	fprintf(fd, "tloaderFlags 0x%xn", m_nt_header->loaderFlags);
	fprintf(fd, "tnumDataDirectories 0x%xn", m_nt_header->numDataDirectories);
	// data directories
	fprintf(fd, "nDataDirectoriesn");
	fprintf(fd, "t;RVAttSizen");
	for(i=0; i< m_nt_header->numDataDirectories; i++)
	{
		fprintf(fd, "t0x%xtt0x%x", 
		m_directories[i].RVA, m_directories[i].size);
		if (i < (sizeof(dataDirNames) / sizeof(char*)))
			fprintf(fd, "tt; %s", dataDirNames[i]);
		else
			fprintf(fd, "tt; !!!UNKNOWN DIRECTORY!!!");
		fprintf(fd, "n");
	}
	// sections
	fprintf(fd, "nSectionsn");
	fprintf(fd, "t;NamettVSizetRVAtSizetOffsettReltLinest#Relt#LinetFlagsn");
	for(i=0; i < m_pe_header->numSections; i++)
	{
		strncpy(buffer, (char*)(m_sections[i].name), 8);
		*(buffer+8) = 0;
		fprintf(fd, "t%st0x%xt0x%xt0x%xt0x%xt0x%xt0x%xt0x%xt0x%xt0x%xtn", 
		buffer,
		m_sections[i].misc.virtualSize,
		m_sections[i].RVA,
		m_sections[i].dataAlignSize,
		m_sections[i].dataOffset,
		m_sections[i].relocationsOffset,
		m_sections[i].lineNumbersOffset,
		m_sections[i].numRelocations,
		m_sections[i].numLineNumbers,
		m_sections[i].flags);
	}
	fclose(fd);
	return (true);
}
// SaveExecutable
//
// - Write the current executable
//==============================================================================================
bool CPEFile::SaveExecutable(char *filename)
{
	FILE *f;
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	// Write the current buffer into a file
	f = fopen(filename, "wb");
	if (!f)
	{
		MessageBox(m_hwnd, CString("Can't open for writing the file ") + CString(filename), "Open error", 0);
		return (false);
	}
	if (fwrite(m_buffer, sizeof(unsigned char), m_size, f) !=
		(size_t)m_size)
	{
		MessageBox(m_hwnd, "Can't write the whole executable", "Open error", 0);
		fclose(f);
		return (false);
	}
	fclose(f);
	return (true);
}
// SavePartialExecutable
//
// - Write the current executable in partial mode
//==============================================================================================
bool CPEFile::SavePartialExecutable(char *filename, DWORD start, DWORD length)
{
	FILE *f;
	start -= m_nt_header->imageBase;
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	// Write the current buffer into a file
	f = fopen(filename, "wb");
	if (!f)
	{
		MessageBox(m_hwnd, CString("Can't open for writing the file ") + CString(filename), "Open error", 0);
		return (false);
	}
	if (fwrite(m_buffer+start, sizeof(unsigned char), length, f) !=
		(size_t)length)
	{
		MessageBox(m_hwnd, "Can't write a part of this executable", "Open error", 0);
		fclose(f);
		return (false);
	}
	fclose(f);
	return (true);
}
// SaveHeaderOnly
//
// - Write the header of the current executable
//==============================================================================================
bool CPEFile::SaveHeaderOnly(char *filename)
{
	FILE *f;
	DWORD dwHeaderSize;
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	// Write the current buffer into a file
	f = fopen(filename, "wb");
	if (!f)
	{
		MessageBox(m_hwnd, CString("Can't open for writing the file ") + CString(filename), "Open error", 0);
		return (false);
	}
	dwHeaderSize = m_dosstub_size +
		sizeof(PEHeader) +
		sizeof(NTOptionalHeader) +
		sizeof(DataDirectory)*m_nt_header->numDataDirectories +
		sizeof(Section)*m_pe_header->numSections;
	// Save the whole headers
	if (fwrite(m_buffer, sizeof(unsigned char), dwHeaderSize, f) !=
		(size_t)dwHeaderSize)
	{
		MessageBox(m_hwnd, "Can't write the whole headers", "Open error", 0);
		fclose(f);
		return (false);
	}
	fclose(f);
	return (true);
}
int CPEFile::FindSectionIndex(DWORD addr)
{
	int i;
	for (i=0; i<m_pe_header->numSections; i++)
	{
		if (addr >= m_sections[i].RVA && 
			addr < m_sections[i].RVA + m_sections[i].misc.virtualSize)
		{
			break;
		}
	}
	if (i<m_pe_header->numSections)
		return (i);
	else
		return (-1);
}
DWORD CPEFile::RVA2Offset(DWORD addr)
{
	int i = FindSectionIndex(addr);
	if (i >= 0)
		return (addr - m_sections[i].RVA + m_sections[i].dataOffset);
	return (NULL);
}
// GetHeader
//
// - Replace the header of buffer into the current
//==============================================================================================
bool CPEFile::FixHeader()
{
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	CPEFile pe_file_header(m_hwnd);
	if (pe_file_header.LoadExecutable(m_filename) &&
		pe_file_header.UpdatePEVars(true))	// FIX RAW=RVA for all sections
	{
		memcpy(m_buffer, pe_file_header.m_buffer, pe_file_header.m_nt_header->headersSize);
		return (true);
	}
	return (false);
}
// Rebuild Import
//
// - Rebuild the import table
//==============================================================================================
bool CPEFile::RebuildImport(void **name_import)
{
	// Executable not yet loaded!!!!
	if (!m_buffer)
	{
		MessageBox(m_hwnd, "Open a file before!!", "Open error", 0);
		return (false);
	}
	char buf[256];
	int  i;
	if ((i=FindSectionIndex(m_directories[ImportDataDirectory].RVA)) >= 0)
	{
		int nb_dll = 0, nn = 0;
		// Cut the import table to 3 parts
		//
		// 1 - Image Import Descriptor
		// 2 - Import Array Table
		// 3 - Function Name
		// FIRST
		IMAGE_IMPORT_DESCRIPTOR *tmp = 
			(IMAGE_IMPORT_DESCRIPTOR*)(m_buffer+m_sections[i].dataOffset+
			m_directories[ImportDataDirectory].RVA-m_sections[i].RVA);
		DWORD *tmp2;
		void *imp1 = tmp;
		void *imp2;
		void *imp3;
		// SECOND
		while (tmp->Name)
		{
			tmp++;
			nb_dll++;
		}
		tmp++;
		imp2 = tmp;
		// THIRD
		tmp2 = (DWORD*)tmp;
		while (nn < nb_dll)
		{
			if (!(*tmp2))
			{
				nn++;
			}
			tmp2++;
		}
		imp3 = tmp2;
		sprintf(buf, "%X %X %X",
			(DWORD)imp1-(DWORD)m_buffer,
			(DWORD)imp2-(DWORD)m_buffer,
			(DWORD)imp3-(DWORD)m_buffer);
		*name_import = imp3;
		MessageBox(m_hwnd, buf, "Ok", 0);
		return (true);
	}
	MessageBox(m_hwnd, "Argh! No section found.!!", "Import error", 0);
	return (false);
}

						