//1: pe.h 文件

/*
 ===== PE structures and types definition =====
*/

 

#ifndef PE_H
#define PE_H

typedef struct
{
  unsigned short cpuType;
  unsigned short numSections;
  unsigned long dateStamp;
  unsigned long symbolTable;
  unsigned long numSymbols;
  unsigned short optionalHeaderSize;
  unsigned short flags;
} PEHeader;


typedef struct
{
  unsigned short magic;
  unsigned char linkerMajor;
  unsigned char linkerMinor;
  unsigned long codeSize;
  unsigned long initDataSize;
  unsigned long uninitDataSize;
  unsigned long entryPoint;
  unsigned long codeBase;
  unsigned long dataBase;
} StdOptionalHeader;


typedef struct
{
  unsigned short magic;
  unsigned char linkerMajor;
  unsigned char linkerMinor;
  unsigned long codeSize;
  unsigned long initDataSize;
  unsigned long uninitDataSize;
  unsigned long entryPoint;
  unsigned long codeBase;
  unsigned long dataBase;

  // extra NT stuff
  unsigned long p_w_picpathBase;
  unsigned long sectionAlign;
  unsigned long fileAlign;
  unsigned short osMajor;
  unsigned short osMinor;
  unsigned short p_w_picpathMajor;
  unsigned short p_w_picpathMinor;
  unsigned short subsystemMajor;
  unsigned short subsystemMinor;
  unsigned long reserved;
  unsigned long p_w_picpathSize;
  unsigned long headersSize;
  unsigned long checksum;
  unsigned short subsystem;
  unsigned short dllFlags;
  unsigned long stackReserveSize;
  unsigned long stackCommitSize;
  unsigned long heapReserveSize;
  unsigned long heapCommitSize;
  unsigned long loaderFlags;
  unsigned long numDataDirectories;
} NTOptionalHeader;

typedef struct
{
  unsigned short magic;
  unsigned char linkerMajor;
  unsigned char linkerMinor;
  unsigned long codeSize;
  unsigned long initDataSize;
  unsigned long uninitDataSize;
  unsigned long entryPoint;
  unsigned long codeBase;
  unsigned long dataBase;
  unsigned long bssBase;
  unsigned long gprMask;
  unsigned long cprMask[4];
  unsigned long gpValue;
} ROMOptionalHeader;

#define NTOptionalHeaderMagic 0x10b
#define ROMOptionalHeaderMagic 0x107

typedef struct
{
  unsigned long RVA;
  unsigned long size;
} DataDirectory;

typedef struct
{
  unsigned char name[8];
  union
  {
    unsigned long virtualSize;
    unsigned long physicalAddress;
  } misc;
  unsigned long RVA;
  unsigned long dataAlignSize;
  unsigned long dataOffset;
  unsigned long relocationsOffset; //}not actually used in PE
  unsigned long lineNumbersOffset; //}
  unsigned short numRelocations;   //}
  unsigned short numLineNumbers;   //}
  unsigned long flags;
} Section;


#define ExportDataDirectory 0
#define ImportDataDirectory 1
#define ResourceDataDirectory 2
#define ExceptionDataDirectory 3
#define SecurityDataDirectory 4
#define BaseRelocDataDirectory 5
#define DebugDataDirectory 6
#define CopyrightDataDirectory 7
#define GlobalPtrdataDirectory 8
#define TlsDataDirectory 9
#define LoadConfigDataDirectory 10

static char *dataDirNames[] = {"Export",
          "Import",
          "Resource",
          "Exception",
          "Security",
          "BaseReloc",
          "Debug",
          "Copyright",
          "GlobalPtr",
          "TlsData",
          "LoadConfig",
          "BoundImport",
          "IAT"};

#endif
 

//2.PEFile.h文件

// PEFile.h: interface for the CPEFile class.
//
//

#if !defined(AFX_PEFILE_H__BC746100_7BF3_11D4_8B4C_DF3F74FBAB11__INCLUDED_)
#define AFX_PEFILE_H__BC746100_7BF3_11D4_8B4C_DF3F74FBAB11__INCLUDED_

#if _MSC_VER > 1000
#pragma once
#endif // _MSC_VER > 1000

#include "pe.h"

class CPEFile 
{
// Public methods
public:
 CPEFile(HWND hwnd = NULL);
 virtual ~CPEFile();

 bool  LoadExecutable(char *filename);
 bool  WriteInfos(char *filename);
 bool  SaveExecutable(char *filename);
 bool  SavePartialExecutable(char *filename, DWORD start, DWORD length);
 bool  SaveHeaderOnly(char *filename);
 bool  FixHeader();
 bool  UpdatePEVars(bool fix_sections = false);
 bool  RebuildImport(void **name_import);
 int   FindSectionIndex(DWORD addr);
 DWORD RVA2Offset(DWORD addr);

// Protected methods
protected:

// Protected datas
public:
 HWND    m_hwnd;      // for messagebox (can be NULL)
 char    m_filename[_MAX_PATH];  // name of executable
 unsigned char  *m_buffer;     // buffer of executable
 long    m_size;      // size of executable

 // pe infos
 unsigned int  m_dosstub_size;
 PEHeader   *m_pe_header;
 StdOptionalHeader *m_std_header;
 NTOptionalHeader *m_nt_header;
 ROMOptionalHeader *m_rom_header;
 DataDirectory  *m_directories;
 Section    *m_sections;
};

#endif // !defined(AFX_PEFILE_H__BC746100_7BF3_11D4_8B4C_DF3F74FBAB11__INCLUDED_)
 

//3.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;
 }
}


// LoadExecutable
//
// - Load the executable in memory and parse the pe header
//==============================================================================================
bool CPEFile::LoadExecutable(char *filename)
{
 bool result;
 FILE *f;

 // 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_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 ====---\n\n");
 fprintf(fd, "\tdosStubSize 0x%x\n", m_dosstub_size);
 fprintf(fd, "\nPEHeader\n");
 fprintf(fd, "\tcpuType 0x%x\n", m_pe_header->cpuType);
 fprintf(fd, "\tnumSections 0x%x\n", m_pe_header->numSections);
 fprintf(fd, "\tdateStamp 0x%x\n", m_pe_header->dateStamp);
 fprintf(fd, "\tsymbolTable 0x%x\n", m_pe_header->symbolTable);
 fprintf(fd, "\tnumSymbols 0x%x\n", m_pe_header->numSymbols);
 fprintf(fd, "\toptionalHeaderSize 0x%x\n", m_pe_header->optionalHeaderSize);
 fprintf(fd, "\tflags 0x%x\n", m_pe_header->flags);

 // pe optional header
 fprintf(fd, "\nNTOptionalHeader\n");
 fprintf(fd, "\tmagic 0x%x\n", m_nt_header->magic);
 fprintf(fd, "\tlinkerMajor 0x%x\n", m_nt_header->linkerMajor);
 fprintf(fd, "\tlinkerMinor 0x%x\n", m_nt_header->linkerMinor);
 fprintf(fd, "\tcodeSize 0x%x\n", m_nt_header->codeSize);
 fprintf(fd, "\tinitDataSize 0x%x\n", m_nt_header->initDataSize);
 fprintf(fd, "\tuninitDataSize 0x%x\n", m_nt_header->uninitDataSize);
 fprintf(fd, "\tentryPoint 0x%x\n", m_nt_header->entryPoint);
 fprintf(fd, "\tcodeBase 0x%x\n", m_nt_header->codeBase);
 fprintf(fd, "\tdataBase 0x%x\n", m_nt_header->dataBase);
 fprintf(fd, "\tp_w_picpathBase 0x%x\n", m_nt_header->p_w_picpathBase);
 fprintf(fd, "\tsectionAlign 0x%x\n", m_nt_header->sectionAlign);
 fprintf(fd, "\tfileAlign 0x%x\n", m_nt_header->fileAlign);
 fprintf(fd, "\tosMajor 0x%x\n", m_nt_header->osMajor);
 fprintf(fd, "\tosMinor 0x%x\n", m_nt_header->osMinor);
 fprintf(fd, "\tp_w_picpathMajor 0x%x\n", m_nt_header->p_w_picpathMajor);
 fprintf(fd, "\tp_w_picpathMinor 0x%x\n", m_nt_header->p_w_picpathMinor);
 fprintf(fd, "\tsubsystemMajor 0x%x\n", m_nt_header->subsystemMajor);
 fprintf(fd, "\tsubsystemMinor 0x%x\n", m_nt_header->subsystemMinor);
 fprintf(fd, "\treserved 0x%x\n", m_nt_header->reserved);
 fprintf(fd, "\tp_w_picpathSize 0x%x\n", m_nt_header->p_w_picpathSize);
 fprintf(fd, "\theadersSize 0x%x\n", m_nt_header->headersSize);
 fprintf(fd, "\tchecksum 0x%x\n", m_nt_header->checksum);
 fprintf(fd, "\tsubsystem 0x%x\n", m_nt_header->subsystem);
 fprintf(fd, "\tdllFlags 0x%x\n", m_nt_header->dllFlags);
 fprintf(fd, "\tstackReserveSize 0x%x\n", m_nt_header->stackReserveSize);
 fprintf(fd, "\tstackCommitSize 0x%x\n", m_nt_header->stackCommitSize);
 fprintf(fd, "\theapReserveSize 0x%x\n", m_nt_header->heapReserveSize);
 fprintf(fd, "\theapCommitSize 0x%x\n", m_nt_header->heapCommitSize);
 fprintf(fd, "\tloaderFlags 0x%x\n", m_nt_header->loaderFlags);
 fprintf(fd, "\tnumDataDirectories 0x%x\n", m_nt_header->numDataDirectories);

 // data directories
 fprintf(fd, "\nDataDirectories\n");
 fprintf(fd, "\t;RVA\t\tSize\n");
 for(i=0; i< m_nt_header->numDataDirectories; i++)
 {
  fprintf(fd, "\t0x%x\t\t0x%x",
  m_directories[i].RVA, m_directories[i].size);

  if (i < (sizeof(dataDirNames) / sizeof(char*)))
   fprintf(fd, "\t\t; %s", dataDirNames[i]);
  else
   fprintf(fd, "\t\t; !!!UNKNOWN DIRECTORY!!!");

  fprintf(fd, "\n");
 }

 // sections
 fprintf(fd, "\nSections\n");
 fprintf(fd, "\t;Name\t\tVSize\tRVA\tSize\tOffset\tRel\tLines\t#Rel\t#Line\tFlags\n");
 for(i=0; i < m_pe_header->numSections; i++)
 {
  strncpy(buffer, (char*)(m_sections[i].name), 8);
  *(buffer+8) = 0;
  fprintf(fd, "\t%s\t0x%x\t0x%x\t0x%x\t0x%x\t0x%x\t0x%x\t0x%x\t0x%x\t0x%x\t\n",
  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->p_w_picpathBase;

 // 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);
}