#include "LPEL_Sections.h"
#include "LPEL_Useful.h"
#include "LPEL_Headers.h"

DWORD DLLEXPORT LPEL_GetSectionRawSize(HANDLE hBinary, PUCHAR pSectionName)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     size_t i;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           if(!strcmp(pSectionHeader[i].Name, pSectionName))
                                              return (DWORD)pSectionHeader[i].SizeOfRawData;

     }
     
     return 0x0;
}
PDWORD DLLEXPORT LPEL_GetSectionRawAddress(HANDLE hBinary, PUCHAR pSectionName, BYTE RelativeAddress)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     size_t i;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           if(!strcmp(pSectionHeader[i].Name, pSectionName))
           {
                                
                  if(RelativeAddress)              
                                 return (PDWORD)pSectionHeader[i].PointerToRawData;
                  else
                                 return (PDWORD)((PUCHAR)hBinary + pSectionHeader[i].PointerToRawData);
           }

     }
     
     return 0x0;
}
PDWORD DLLEXPORT LPEL_GetSectionVirtualAddress(HANDLE hBinary, PUCHAR pSectionName, BYTE RelativeAddress)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;
     
     size_t i;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     pOptionalHeader = LPEL_GetOptionalHeader(hBinary);
     
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           if(!strcmp(pSectionHeader[i].Name, pSectionName))
           {
                    if(RelativeAddress)                          
                                return (PDWORD)pSectionHeader[i].VirtualAddress;
                    else
                                return (PDWORD)((PUCHAR)pOptionalHeader->ImageBase + pSectionHeader[i].VirtualAddress);
           }
                                              
     }
     
     return 0x0;
}
     
BYTE DLLEXPORT LPEL_WriteInSection(HANDLE hBinary, PUCHAR pSectionName, PUCHAR pBuf, UINT size, DWORD dwOffset)
{
     PDWORD pSectionAddress = NULL;
     size_t i;
     
     pSectionAddress = (PDWORD)((PUCHAR)LPEL_GetSectionRawAddress(hBinary, pSectionName, 0)+dwOffset);
     
     if(pSectionAddress == 0x0)
                        return 0x0;
     
     //printf("\nWriting 0x%x bytes at 0x%x", size, pSectionAddress);
     memcpy(pSectionAddress, pBuf, size);
     return 0x1;
}
VOID DLLEXPORT LPEL_ListSections(HANDLE hBinary)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     size_t i;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     
     printf("\n[+]Sections");
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           printf("\n    %s", pSectionHeader[i].Name);
           //printf("\n           VirtualSize : 0x%x", pSectionHeader[i].Misc.VirtualSize);
           //printf("\n           Virtual Address : 0x%x", pSectionHeader[i].VirtualAddress);
           //printf("\n           SizeOfRawData : 0x%x", pSectionHeader[i].SizeOfRawData);
           //printf("\n           PointerToRawData : 0x%x", pSectionHeader[i].PointerToRawData);
           //printf("\n           PointerToRelocations : 0x%x", pSectionHeader[i].PointerToRelocations);
           //printf("\n           PointerToLinenumbers : 0x%x", pSectionHeader[i].PointerToLinenumbers);
           //printf("\n           NumberOfRelocations : 0x%x", pSectionHeader[i].NumberOfRelocations);
           //printf("\n           NumberOfLinenumbers : 0x%x", pSectionHeader[i].NumberOfLinenumbers);
           //printf("\n           Attributes : 0x%x", pSectionHeader[i].Characteristics);
  
     }
}

PUCHAR DLLEXPORT LPEL_GetRawAddressSectionOwner(HANDLE hBinary, DWORD Addr)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;

     size_t i;

     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     
     for(i = 0; i <= pCOFFHeader->NumberOfSections; i++)
     {
           //On regarde si l'adresse est bien dans cette section
           if((PDWORD)Addr >= (PDWORD)pSectionHeader[i].PointerToRawData && (PDWORD)Addr <= (PDWORD)((PUCHAR)pSectionHeader[i].PointerToRawData + pSectionHeader[i].SizeOfRawData))
                                          return pSectionHeader[i].Name;
                                          
     }
     
     return 0x0;
}
DWORD DLLEXPORT LPEL_GetSectionFlag(HANDLE hBinary, PUCHAR pSectionName)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     
     size_t i;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           if(!strcmp(pSectionHeader[i].Name, pSectionName))
                                              return pSectionHeader[i].Characteristics; 
     }
     
     return 0x0;
}
DWORD DLLEXPORT LPEL_SetSectionFlag(HANDLE hBinary, PUCHAR pSectionName, DWORD dwFlag)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     DWORD dwOldFlag;
     
     size_t i;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           if(!strcmp(pSectionHeader[i].Name, pSectionName))
           {
                                              dwOldFlag = pSectionHeader[i].Characteristics;
                                              *(&pSectionHeader[i].Characteristics) = dwFlag;
                                              return dwOldFlag; 
           }
     }
     
     return 0x0;
}
VOID DLLEXPORT LPEL_MemCopySection(HANDLE hBinary, PUCHAR pMem, PIMAGE_SECTION_HEADER pSection)
{
     memcpy(pMem, (PDWORD)((PUCHAR)pSection->PointerToRawData + (ULONG)hBinary), pSection->SizeOfRawData);
     
     return;
}
VOID DLLEXPORT LPEL_MemCopySections(HANDLE hBinary, PUCHAR pMem, PDWORD pdwBytesCopied)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;

     size_t i, j;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     
     LPEL_MemCopySection(hBinary, pMem, (PIMAGE_SECTION_HEADER)pSectionHeader);
     (PUCHAR)*pdwBytesCopied += pSectionHeader[0].SizeOfRawData;
     
     //printf("\n[+]Copying %s (size : 0x%x...", pSectionHeader[0].Name, *pdwBytesCopied);
     
     for(i = 1; i < pCOFFHeader->NumberOfSections; i++)
     {
           LPEL_MemCopySection(hBinary, (PUCHAR)((PUCHAR)pMem + *pdwBytesCopied), (PIMAGE_SECTION_HEADER)pSectionHeader+i);
          // printf("\n[+]Copying %s (size : 0x%x...", pSectionHeader[i].Name, *pdwBytesCopied);
           
           (PUCHAR)*pdwBytesCopied += pSectionHeader[i].SizeOfRawData;
     }
     
           
     return;
}
VOID DLLEXPORT LPEL_EraseSection(HANDLE hBinary, PIMAGE_SECTION_HEADER pSection)
{
     //printf("\n[+]Erasing section at 0x%x, size : 0x%x", pSection->PointerToRawData, pSection->SizeOfRawData);
     if((PUCHAR)pSection->PointerToRawData != 0x00)
           memset(((PUCHAR)pSection->PointerToRawData + (ULONG)hBinary), 0x00, pSection->SizeOfRawData);
     
     return;
}
VOID DLLEXPORT LPEL_EraseSections(HANDLE hBinary)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;
     
     DWORD dwSectionAlignment, dwFileAlignment;
     //DWORD dwBytesErased;
     
     size_t i;
     
     //dwBytesErased = 0x0;
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     pOptionalHeader = LPEL_GetOptionalHeader(hBinary);
     
     dwSectionAlignment = pOptionalHeader->SectionAlignment;
     dwFileAlignment = pOptionalHeader->FileAlignment;
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           LPEL_EraseSection(hBinary, (PIMAGE_SECTION_HEADER)pSectionHeader+i);
           //dwBytesErased += pSectionHeader[i].Misc.VirtualSize;
     }
     
     //Updating SizeOfImage      
     /*printf("\n             SizeOfImage updated from 0x%x to", pOptionalHeader->SizeOfImage);
     pOptionalHeader->SizeOfImage = LPEL_GetAlignment(pOptionalHeader->SizeOfImage-dwBytesErased, dwSectionAlignment);
     printf(" 0x%x", pOptionalHeader->SizeOfImage);*/
     
     return;
}

VOID DLLEXPORT LPEL_EraseSectionTable(HANDLE hBinary)
{
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;
     
     DWORD dwSectionAlignment, dwFileAlignment;
     
     size_t i;
     
     //printf("\n[+]Erasing section table...");
     
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     pOptionalHeader = LPEL_GetOptionalHeader(hBinary);
     
     dwSectionAlignment = pOptionalHeader->SectionAlignment;
     dwFileAlignment = pOptionalHeader->FileAlignment;
     
     for(i = 0; i < pCOFFHeader->NumberOfSections; i++)
     {
           memset((PIMAGE_SECTION_HEADER)pSectionHeader+i, 0x00, sizeof(IMAGE_SECTION_HEADER));
     }
     *(&pCOFFHeader->NumberOfSections) -= i;
     //(PUCHAR)*(&pOptionalHeader->SizeOfImage) -= sizeof(IMAGE_SECTION_HEADER)*i;
     
     //Updating SizeOfHeader
     //printf("\n             SizeOfHeader updated from 0x%x to", pOptionalHeader->SizeOfHeaders);
     pOptionalHeader->SizeOfHeaders = LPEL_GetAlignment(pOptionalHeader->SizeOfHeaders - sizeof(IMAGE_SECTION_HEADER)*i, dwFileAlignment);
     //printf(" 0x%x", pOptionalHeader->SizeOfHeaders);

     //Updating SizeOfImage
     //printf("\n             SizeOfImage updated from 0x%x to", pOptionalHeader->SizeOfImage);
     pOptionalHeader->SizeOfImage = LPEL_GetAlignment(pOptionalHeader->SizeOfImage-sizeof(IMAGE_SECTION_HEADER)*i, dwSectionAlignment);
    // printf(" 0x%x", pOptionalHeader->SizeOfImage);
     
     return;
}
           
VOID DLLEXPORT LPEL_CreateSection(HANDLE hBinary, PUCHAR pSectionName, PDWORD pdwVirtualAddress, DWORD dwVirtualSize, PDWORD pdwRawAddress, DWORD dwRawSize, DWORD dwSectionFlag)
{
     PIMAGE_SECTION_HEADER pNewSectionHeader = NULL;
     PIMAGE_SECTION_HEADER pSectionHeader = NULL;
     PIMAGE_FILE_HEADER pCOFFHeader = NULL;
     PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;
     
     DWORD dwSectionSize, dwSectionAlignment, dwFileAlignment;
        
     pCOFFHeader = LPEL_GetCOFFHeader(hBinary);
     pOptionalHeader = LPEL_GetOptionalHeader(hBinary);
     pSectionHeader = LPEL_GetSectionHeader(hBinary);
     
     
     dwSectionAlignment = pOptionalHeader->SectionAlignment;
     dwFileAlignment = pOptionalHeader->FileAlignment;
     
     
     if(pCOFFHeader->NumberOfSections == 0x0)
     {
          pNewSectionHeader = (PIMAGE_SECTION_HEADER)((PUCHAR)pOptionalHeader + pCOFFHeader->SizeOfOptionalHeader);
     }
     else
     {
         pNewSectionHeader = (PIMAGE_SECTION_HEADER) ((PUCHAR)(&pSectionHeader[pCOFFHeader->NumberOfSections-1].Characteristics) + 0x4);
     }
     
     /*printf("\n[+]Adding section %s at address 0x%x ...", pSectionName, pNewSectionHeader);
     printf("\n          Raw Offset : 0x%x",  GetAlignment(pdwRawAddress, dwFileAlignment));
     printf("\n          Size : 0x%x", GetAlignment(dwRawSize, dwFileAlignment));*/
          
     memcpy(*(&pNewSectionHeader->Name), pSectionName, strlen(pSectionName));

     pNewSectionHeader->VirtualAddress = LPEL_GetAlignment(pdwVirtualAddress, dwSectionAlignment);
     pNewSectionHeader->Misc.VirtualSize = LPEL_GetAlignment(dwVirtualSize, dwSectionAlignment);
     pNewSectionHeader->SizeOfRawData = LPEL_GetAlignment(dwRawSize, dwFileAlignment);
     pNewSectionHeader->PointerToRawData = LPEL_GetAlignment(pdwRawAddress, dwFileAlignment);//GetAlignment(400, dwFileAlignment);
     pNewSectionHeader->Characteristics = dwSectionFlag;
     pNewSectionHeader->PointerToRelocations = 0x0;
     pNewSectionHeader->PointerToLinenumbers = 0x0;
     pNewSectionHeader->NumberOfRelocations = 0x0;
     pNewSectionHeader->NumberOfLinenumbers = 0x0;
     
     /*Champs à ne pas oublier*/
     pCOFFHeader->NumberOfSections += 0x1;
     
     //printf("\n          SizeOfImage 0x%x + 0x%x =", pOptionalHeader->SizeOfImage, dwVirtualSize);
     pOptionalHeader->SizeOfImage = LPEL_GetAlignment(pOptionalHeader->SizeOfImage+dwVirtualSize, dwSectionAlignment);
     //printf(" 0x%x", pOptionalHeader->SizeOfImage);
     pOptionalHeader->SizeOfHeaders = LPEL_GetAlignment(pOptionalHeader->SizeOfHeaders + sizeof(IMAGE_SECTION_HEADER), dwFileAlignment);

     return;
}