#include "main.h"

#define path "test.esc.exe"
#define path_dump "Dump.exe"

int			main(int argc, char **argv)
{
  HANDLE		hFile_dump;
  HANDLE		hFile, hMapp;
  char			*mapping = NULL;
  PIMAGE_DOS_HEADER	pDosHeader;
  PIMAGE_NT_HEADERS	pPE;
  PIMAGE_SECTION_HEADER	pSection;
  int			i;
  STARTUPINFOA		StartupInfo;
  PROCESS_INFORMATION	ProcessInfo;
  DEBUG_EVENT		debug_event = {0};
  CONTEXT		Context;
  DWORD			new_oep;
  PCHAR			Dump;
  DWORD			size;
  DWORD			written;

  hFile = CreateFileA(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
  if (hFile == INVALID_HANDLE_VALUE)
    {
      printf("[-] Echec lors de l'ouverture du fichier\n");
      return (0);
    }

  hMapp = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
  mapping = (char*)MapViewOfFile(hMapp, FILE_MAP_READ, 0, 0, 0);
  if (mapping == NULL)
    {
      printf("[-] Echec lors du mapping du fichier\n");
      goto error;
    }
  pDosHeader = (PIMAGE_DOS_HEADER)mapping;
  if (!check_msdos_header(pDosHeader))
    {
      printf("[-] Dos header non valide\n");
      goto error;
    }
  pPE = (PIMAGE_NT_HEADERS)(pDosHeader->e_lfanew + mapping);
  if (!check_pe_header(pPE))
    {
      printf("[-] Pe Header non valide\n");
      goto error;
    }
  pSection = (PIMAGE_SECTION_HEADER)((PCHAR)pPE + sizeof(IMAGE_FILE_HEADER) +
  				     pPE->FileHeader.SizeOfOptionalHeader + sizeof(DWORD));
  if (!check_section_name(pSection))
    {
      printf("[-] La signature .@ n'a pas ete trouve dans le nom des sections\n");
      goto error;
    }

  if (!check_signature(pPE, pSection, mapping))
    {
      printf("Bad Signature\n");
      goto error;
    }

  ZeroMemory(&StartupInfo, sizeof(StartupInfo));
  StartupInfo.cb = sizeof(StartupInfo);
  ZeroMemory(&ProcessInfo, sizeof(ProcessInfo));
  CreateProcessA(path, NULL, NULL, NULL, FALSE, DEBUG_ONLY_THIS_PROCESS | DEBUG_PROCESS,
  		 NULL, NULL, &StartupInfo, &ProcessInfo);

  disable_isdebugerpresent(ProcessInfo.hProcess);
  Context.ContextFlags = CONTEXT_ALL;
  BreakAndNbStep(pPE->OptionalHeader.AddressOfEntryPoint + pPE->OptionalHeader.ImageBase, ProcessInfo.hThread, 3, &Context, BreakOnExec, &debug_event);
  GetThreadContext(ProcessInfo.hThread,&Context);
  BreakAndNbStep(Context.Esp, ProcessInfo.hThread, 3, &Context, BreakOnAccess, &debug_event);
  new_oep = Context.Eip;
  Dump = DumpProcess(ProcessInfo.hProcess, (char*)pPE->OptionalHeader.ImageBase, &size);
  if (!Dump)
    {
      printf("[-] DumpProcess() failed\n");
      goto error;
    }
  hFile_dump = CreateFileA(path_dump, GENERIC_WRITE, FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, 0, NULL);
  if (hFile_dump == INVALID_HANDLE_VALUE)
    {
      printf("[-] Echec lors de l'ouverture du fichier de dump\n");
      goto error;
    }
  RebuildDump(Dump, new_oep);
  WriteFile(hFile_dump, Dump, size, (LPDWORD)&written, NULL);
  CloseHandle(hFile_dump);
  printf("Use ImportRec with OEP : %x\n", new_oep - pPE->OptionalHeader.ImageBase);
  system("pause");
  ContinueDebugEvent(debug_event.dwProcessId,debug_event.dwThreadId,DBG_CONTINUE);
  while(WaitForDebugEvent(&debug_event,INFINITE))
    {
      if (debug_event.dwDebugEventCode == EXCEPTION_DEBUG_EVENT)
  	{
  	  if (debug_event.u.Exception.ExceptionRecord.ExceptionCode == STATUS_BREAKPOINT)
  	    {
  	      ContinueDebugEvent(debug_event.dwProcessId,
  				 debug_event.dwThreadId,
  				 DBG_CONTINUE);
  	    }
  	  else if (debug_event.u.Exception.ExceptionRecord.ExceptionCode == EXCEPTION_SINGLE_STEP)
	    ContinueDebugEvent(debug_event.dwProcessId, debug_event.dwThreadId, DBG_CONTINUE);
  	  else
  	    ContinueDebugEvent(debug_event.dwProcessId, debug_event.dwThreadId,
			       DBG_EXCEPTION_NOT_HANDLED);
  	}
      else if (debug_event.dwDebugEventCode == EXIT_PROCESS_DEBUG_EVENT)
      	{
      	  break;
      	}
      else
  	ContinueDebugEvent(debug_event.dwProcessId, debug_event.dwThreadId, DBG_CONTINUE);
    }
  return (0);

 error:
  CloseHandle(hMapp);
  CloseHandle(hFile);
  return (1);
}

int		check_signature(PIMAGE_NT_HEADERS pPE, PIMAGE_SECTION_HEADER pSection, char *mapping)
{
  ULONG		oep;
  ULONG		oep_offset;
  DWORD		i;

  oep = pPE->OptionalHeader.AddressOfEntryPoint;
  pSection = (PIMAGE_SECTION_HEADER)((PCHAR)pPE + sizeof(IMAGE_FILE_HEADER) +
  				     pPE->FileHeader.SizeOfOptionalHeader + sizeof(DWORD));
  do
    {
      if (pSection->VirtualAddress >= oep && oep <= pSection->VirtualAddress +
	  pSection->Misc.VirtualSize)
  	break;
      pSection = (PIMAGE_SECTION_HEADER)((PCHAR)pSection +
  					 sizeof(IMAGE_SECTION_HEADER));
    }while(1);

  oep_offset = oep - pSection->VirtualAddress + pSection->PointerToRawData;
  for (i = 0; i < strlen(sig); i++)
    if (sig[i] != *(oep_offset + mapping + i))
      return (0);
  return (1);
}

DWORD				disable_isdebugerpresent(HANDLE hProcess)
{
  lpfNtQueryInformationProcess	NtQueryInformationProcess = NULL;
  PEB				Peb = {0};
  PROCESS_BASIC_INFORMATION	pbi = {0};
  NTSTATUS			status;
  DWORD				dwLength;
  HMODULE			hModule = NULL;
  DWORD				dwBytesRead = 0;

  hModule = GetModuleHandleA("ntdll.dll");
  if (hModule == NULL)
    {
      printf("[-] Failed GetModuleHandleA\n");
      return (0);
    }
  NtQueryInformationProcess = (lpfNtQueryInformationProcess)
    GetProcAddress(hModule, "NtQueryInformationProcess");
  FreeLibrary(hModule);
  if (NtQueryInformationProcess == NULL)
    {
      printf("[-] GetProcAddress\n");
      return (0);
    }
  status = NtQueryInformationProcess(hProcess, ProcessBasicInformation, 
				     (PVOID)&pbi, sizeof(PROCESS_BASIC_INFORMATION), &dwLength);
  if (!ReadProcessMemory(hProcess, pbi.PebBaseAddress, &Peb, 
			 sizeof(PEB), &dwBytesRead) || dwBytesRead < sizeof(PEB))
    {
      printf("[-] Error ReadProcessMemory\n");
      return (0);
    }
  Peb.BeingDebugged = 0;
  WriteProcessMemory(hProcess, pbi.PebBaseAddress, &Peb, sizeof(PEB), &dwBytesRead);
  return (1);
}


int check_msdos_header(PIMAGE_DOS_HEADER pDosHeader)
{
  if (pDosHeader->e_magic != 'ZM')
    return (0);
  return (1);
}

int check_pe_header(PIMAGE_NT_HEADERS pPE)
{
  if (pPE->Signature != 'EP')
    return (0);
  return (1);
}

int check_section_name(PIMAGE_SECTION_HEADER	pSection)
{
  do
    {
      if (pSection->VirtualAddress == 0)
	return (0);
      if (!strcmp(pSection->Name, ".@"))
  	break;
      pSection = (PIMAGE_SECTION_HEADER)((PCHAR)pSection +
  					 sizeof(IMAGE_SECTION_HEADER));
    }while(1);
  return (1);
}

void				RebuildDump(char *Dump, DWORD new_oep)
{
  PIMAGE_DOS_HEADER		pDosHeader;
  PIMAGE_NT_HEADERS		pPE;
  PIMAGE_SECTION_HEADER		pSection;
  PIMAGE_IMPORT_DESCRIPTOR	IT;


  pDosHeader = (PIMAGE_DOS_HEADER)Dump;
  pPE = (PIMAGE_NT_HEADERS)(Dump+pDosHeader->e_lfanew);
  pPE->OptionalHeader.AddressOfEntryPoint = new_oep - pPE->OptionalHeader.ImageBase;

  pSection = (PIMAGE_SECTION_HEADER)((PCHAR)pPE + sizeof(IMAGE_FILE_HEADER)
				     + pPE->FileHeader.SizeOfOptionalHeader + sizeof(DWORD));
  while (pSection->VirtualAddress != 0)
    {
      pSection->SizeOfRawData = pSection->Misc.VirtualSize;
      pSection->PointerToRawData = pSection->VirtualAddress;
      pSection++;
    }
}

DWORD	BreakAndNbStep(DWORD addr, HANDLE hThread, int nb_step, CONTEXT *Context, DWORD Flag,
		       DEBUG_EVENT *debug_event)
{
   GetThreadContext(hThread,Context);
   Context->Dr0 = addr;
   Context->Dr6 = 0;
   Context->Dr7 |= SetDR7Flag(OneByteLength, Flag, GlobalFlag | LocalFlag, 0);
   SetThreadContext(hThread, Context);
   ContinueDebugEvent(debug_event->dwProcessId, debug_event->dwThreadId, DBG_CONTINUE);
   while (WaitForDebugEvent(debug_event, INFINITE))
     {
       if (debug_event->dwDebugEventCode == EXCEPTION_DEBUG_EVENT)
	 {
	   if (debug_event->u.Exception.ExceptionRecord.ExceptionCode == STATUS_BREAKPOINT)
	     {
	       ContinueDebugEvent(debug_event->dwProcessId,
				  debug_event->dwThreadId,
				  DBG_CONTINUE);
	     }
	   else if (debug_event->u.Exception.ExceptionRecord.ExceptionCode == EXCEPTION_SINGLE_STEP)
	     {
	       GetThreadContext(hThread,Context);
	       nb_step--;
	       if (!nb_step)
		 return (1);
	       Context->EFlags |= 0x100;
	       Context->Dr0 = 0;
	       Context->Dr7 = 0;
	       SetThreadContext(hThread, Context);
	       ContinueDebugEvent(debug_event->dwProcessId, debug_event->dwThreadId, DBG_CONTINUE);
	     }
	   else if (debug_event->dwDebugEventCode == EXIT_PROCESS_DEBUG_EVENT)
	     {
	       break;
	     }
	   else
	     ContinueDebugEvent(debug_event->dwProcessId, debug_event->dwThreadId,
				DBG_EXCEPTION_NOT_HANDLED);
	 }
       else
	 ContinueDebugEvent(debug_event->dwProcessId, debug_event->dwThreadId, DBG_CONTINUE);
     }
   return (0);
}

PCHAR			DumpProcess(HANDLE hProcess,char *BaseAddress, DWORD *size)
{
  DWORD			AllocSize = 0;
  char			*Dump = NULL;
  IMAGE_DOS_HEADER	DosHeader;
  IMAGE_NT_HEADERS	PE;
  PIMAGE_SECTION_HEADER pSectionHeaders;
  DWORD			i = 0;
  DWORD			SectionHeadersSize, NumberOfBytesRead;


  ReadProcessMemory(hProcess, BaseAddress, &DosHeader, sizeof(IMAGE_DOS_HEADER), 
		    &NumberOfBytesRead);
  if (NumberOfBytesRead != sizeof(IMAGE_DOS_HEADER))
    return (0);
  ReadProcessMemory(hProcess, BaseAddress + DosHeader.e_lfanew, &PE, sizeof(IMAGE_NT_HEADERS), 
		    &NumberOfBytesRead);
  if (NumberOfBytesRead != sizeof(IMAGE_NT_HEADERS))
    return (0);
  SectionHeadersSize = sizeof(IMAGE_NT_HEADERS) * PE.FileHeader.NumberOfSections;
  pSectionHeaders = (PIMAGE_SECTION_HEADER)malloc(SectionHeadersSize);
  ReadProcessMemory(hProcess, BaseAddress + DosHeader.e_lfanew + sizeof(IMAGE_FILE_HEADER) 
		    + PE.FileHeader.SizeOfOptionalHeader +
		    sizeof(DWORD), pSectionHeaders, SectionHeadersSize, &NumberOfBytesRead);
  if (NumberOfBytesRead != SectionHeadersSize)
    {
      free(pSectionHeaders);
      return (0);
    }
  while ((pSectionHeaders + i + 1)->VirtualAddress != 0)
    i++;
  AllocSize = (pSectionHeaders + i)->VirtualAddress + (pSectionHeaders + i)->Misc.VirtualSize;
  free(pSectionHeaders);
  Dump = (char*)VirtualAlloc(NULL, AllocSize, MEM_COMMIT |
			     MEM_RESERVE, PAGE_READWRITE);
  ReadProcessMemory(hProcess,BaseAddress,Dump,AllocSize,&NumberOfBytesRead);
  if (NumberOfBytesRead != AllocSize)
    return (0);
  *size = AllocSize;
  return (Dump);
}