#include<fstream>
#include<iostream>
#include<windows.h>
#include <exception>
#include <string>
using namespace std;

BOOL  IsPeFile(LPVOID  ImageBase)   //判断是否是PE文件结构
{
 PIMAGE_DOS_HEADER  pDosHeader = NULL;
 PIMAGE_NT_HEADERS  pNtHeader  = NULL;
 if(!ImageBase){
     return FALSE;
 }
 pDosHeader = (PIMAGE_DOS_HEADER) ImageBase;
 if(pDosHeader->e_magic != IMAGE_DOS_SIGNATURE){
     return FALSE;
 }
  
 pNtHeader = (PIMAGE_NT_HEADERS)((ULONG_PTR) pDosHeader + pDosHeader -> e_lfanew);
 if(pNtHeader -> Signature != IMAGE_NT_SIGNATURE ){
    return FALSE;
 }
 return TRUE;
}


PIMAGE_NT_HEADERS  GetNtHeader(LPVOID  ImageBase) //获取NT结构指针
{
 PIMAGE_DOS_HEADER  pDosHeader = NULL;
 PIMAGE_NT_HEADERS  pNtHeader  = NULL;
 if(!IsPeFile(ImageBase))
  return  NULL;
 pDosHeader = (PIMAGE_DOS_HEADER)ImageBase;
 pNtHeader = (PIMAGE_NT_HEADERS)((ULONG_PTR)pDosHeader+pDosHeader->e_lfanew);
 return    pNtHeader;
}


PIMAGE_FILE_HEADER  WINAPI  GetFileHeader(LPVOID  Imagebase)
{
 PIMAGE_FILE_HEADER  pFileHeader;
 PIMAGE_NT_HEADERS  pNtHeader = NULL;
 pNtHeader = GetNtHeader(Imagebase);
 if(!pNtHeader)
  return  NULL;
 pFileHeader = & pNtHeader->FileHeader; 
 
 return  pFileHeader;
}

PIMAGE_OPTIONAL_HEADER  GetOptionalHeader(LPVOID  ImageBase)
{
 PIMAGE_OPTIONAL_HEADER  pOptionHeader = NULL;
 PIMAGE_NT_HEADERS  pNtHeader = NULL;
 pNtHeader = GetNtHeader(ImageBase);
 if(!pNtHeader)
  return  NULL;
 pOptionHeader = & pNtHeader->OptionalHeader;
 return  pOptionHeader;
}

void ChangeDwordToString(DWORD d, char *cBuffer)
{
 unsigned char waddress[4] = {0};

 cBuffer[3] = (char)(d>>24)&0xFF;
 cBuffer[2] = (char)(d>>16)&0xFF;
 cBuffer[1] = (char)(d>>8)&0xFF;
 cBuffer[0] = (char)(d)&0xFF;
}

VOID HandleSessionTable(LPVOID file,LPVOID base)
{
 // 后面出现的5都是因为 JMP DWORD 这个汇编指令占5个字符

 int codeslength = 5;
 
 IMAGE_NT_HEADERS *nthead=GetNtHeader(base);
 IMAGE_SECTION_HEADER *sessionhead=(IMAGE_SECTION_HEADER*)((ULONG_PTR)nthead+sizeof(IMAGE_NT_HEADERS));
 if(sessionhead->VirtualAddress==NULL)
  return;
 DWORD sessionnum=nthead->FileHeader.NumberOfSections;
 IMAGE_SECTION_HEADER *p=sessionhead;
 DWORD sFileSize=GetFileSize(base,NULL);
 for(int i=0;i<sessionnum;i++)
 {
  cout<<(char*)p->Name<<" " <<(int)p->SizeOfRawData-(int)p->Misc.VirtualSize<<endl;
  if((int)p->SizeOfRawData-(int)p->Misc.VirtualSize>5 && (p->Characteristics&IMAGE_SCN_MEM_EXECUTE))
  {
   DWORD iWrited = 0;
   DWORD vb=p->VirtualAddress+p->Misc.VirtualSize;
   DWORD fos=p->PointerToRawData+p->Misc.VirtualSize;
   p->Misc.VirtualSize += codeslength;
   // 修改了当前节的大小，写回去，因为有对齐，不用改option header的总大小
   SetFilePointer(file,(ULONG_PTR)nthead + 248 + sizeof(IMAGE_SECTION_HEADER) * i,NULL,FILE_BEGIN);
   WriteFile(file, p, sizeof(IMAGE_SECTION_HEADER),&iWrited,0);

   char buf[4] = {0};
   char cmd[1] = {(char)0xE9};
   ChangeDwordToString(nthead->OptionalHeader.AddressOfEntryPoint, buf);
   SetFilePointer(file,fos,NULL,FILE_BEGIN);
   cout<<"injectpos: 0x"<<hex<<fos<<endl;
   
    if(!WriteFile(file,cmd,1,&iWrited,0)){
       TCHAR  *buffer;
    ::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,NULL,GetLastError(),0,( LPTSTR )&buffer,0,NULL );
    MessageBox(0,buffer,"ok",0);
   }
   

   if(!WriteFile(file,buf,4,&iWrited,0)){
       TCHAR  *buffer;
    ::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,NULL,GetLastError(),0,( LPTSTR )&buffer,0,NULL );
    MessageBox(0,buffer,"ok",0);
   }

   


   DWORD oldentry = nthead->OptionalHeader.AddressOfEntryPoint;

   cout<<"oldentry: 0x"<<hex<<oldentry<<endl;
   
   DWORD newentry = vb;
   cout<<"newentry: 0x"<<hex<<newentry<<endl;
   DWORD writesize=0;

   ChangeDwordToString(newentry, buf);
   UINT fileoffsettoentrypos = ((PIMAGE_DOS_HEADER)base)->e_lfanew + 40;
   SetFilePointer(file,fileoffsettoentrypos,NULL,FILE_BEGIN);
   if(!WriteFile(file,buf, 4 , &writesize,0) )
   {
    TCHAR  *buffer;
    ::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,NULL,GetLastError(),0,( LPTSTR )&buffer,0,NULL );
    MessageBox(0,buffer,"ok",0);
   }
   cout<<"success"<<endl;
   break;
  }
  p++;
 }
}

int main(){
 HANDLE hFile = CreateFile("test.exe",           // open  pe file 
  GENERIC_READ|GENERIC_WRITE,              // open for reading 
  NULL,           // share for reading 
  NULL,                      // no security 
  OPEN_EXISTING,             // existing file only 
  FILE_ATTRIBUTE_NORMAL,   // normal file 
  NULL);                     // no attr. template 
 HANDLE hFileMap = CreateFileMapping(hFile,NULL,PAGE_READWRITE,0,0,NULL);
 if(!hFileMap){
  TCHAR  *buffer ;
  ::FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,NULL,GetLastError(),0,( LPTSTR )&buffer,0,NULL );
  MessageBox(0,buffer,"ok",0);
 }
 LPVOID  lpMemory = MapViewOfFile(hFileMap,FILE_MAP_READ|FILE_MAP_WRITE ,NULL,NULL,NULL);

 if(IsPeFile(lpMemory)){
  //AnalyzeNTHEADER(lpMemory);
  cout<<"yes"<<endl;
  IMAGE_NT_HEADERS *nthead=GetNtHeader(lpMemory);
  IMAGE_OPTIONAL_HEADER *image=GetOptionalHeader(lpMemory);
  cout<<"DataDirectory num:"<<image->NumberOfRvaAndSizes<<endl;
  HandleSessionTable(hFile,lpMemory);
 }else{
     cout<<"no"<<endl;
 }
 UnmapViewOfFile(lpMemory);
 CloseHandle(hFileMap);
 CloseHandle(hFile);

 system("pause");

 return 0;
}