#include "LPEL_Directories.h"
#include "LPEL_Headers.h"

VOID DLLEXPORT LPEL_CopyImportTable(HANDLE hBinary, PUCHAR pMem, PDWORD pdwBytesCopied)
{
     PIMAGE_IMPORT_DESCRIPTOR pImportTable = NULL;
     
     pImportTable = LPEL_GetImportTable(hBinary);
     
     //printf("\n[+] IAT : %x", pImportTable->Name);
     *pdwBytesCopied = LPEL_GetSizeOfDirectory(hBinary, IMAGE_DIRECTORY_ENTRY_IMPORT);
     memcpy(pMem, pImportTable, *pdwBytesCopied);
     
     return;
     
}
VOID DLLEXPORT LPEL_EraseImportTable(HANDLE hBinary)
{
     PDWORD pStartAddress;
     DWORD dwSize;
     
     pStartAddress = (PDWORD) LPEL_GetImportTable(hBinary);

     dwSize = LPEL_GetSizeOfDirectory(hBinary, IMAGE_DIRECTORY_ENTRY_IMPORT);
     //printf("\n[+]Erasing ImportTable (0x%x bytes) at 0x%x ...", dwSize, pStartAddress);
     
     
     memset(pStartAddress, 0x0, dwSize);
     
     return;
}   
     
PUCHAR DLLEXPORT LPEL_BuildImportTable(HANDLE hBinary, PDWORD pImportTableAddress, PUCHAR **Dlls, PUCHAR **Functions, PDWORD pdwSizeOfImportTable)
{
     size_t i, j, k;
     
     PUCHAR pImportTable;
     
     PDWORD pdwStartRVA;
     DWORD dwSizeOfDllNames;
     DWORD dwSizeOfFunctionsNames;
     DWORD dwSizeOfFunctionsHints;
     DWORD dwSizeOfFirstThunks;
     DWORD dwSizeOfImportDescriptors;
     DWORD dwFunctionOffset;
     
     
     IMAGE_IMPORT_DESCRIPTOR ImportDesc;
     IMAGE_IMPORT_BY_NAME ImportName;
     
     //printf("\n\n[+] *** BUILDING IMPORT TABLE  ***");
     
     /* Initialisation */
     
     i=0;
     j=0;
     k=0;
     
     dwSizeOfDllNames = 0;
     dwSizeOfFunctionsNames = 0;
     dwSizeOfFunctionsHints = 0;
     dwSizeOfFirstThunks = 0;
     dwSizeOfImportDescriptors = 0;
     *pdwSizeOfImportTable = 0;
     dwFunctionOffset = 0;
     
     while(Dlls[i] != 0x0)
     {
         //printf("\n   DLL : %s", Dlls[i]);
         
         //Increment size
         dwSizeOfDllNames += strlen(Dlls[i])+1; //include null byte
         
         
         while(Functions[j] != 0x0)
         {
                            //printf("\n              Function : %s", Functions[j]);
                            //Increment size
                            dwSizeOfFunctionsNames += strlen(Functions[i])+1; //include null byte
                            j++;
         }
         i++;
     }
     
     // i-1 = Number of Dlls
     // j-1 = Number of functions
     
     //Increment size
     dwSizeOfFunctionsNames += sizeof(DWORD);
     dwSizeOfFunctionsHints = (j)*sizeof(WORD);
     dwSizeOfImportDescriptors += (i+1)*sizeof(IMAGE_IMPORT_DESCRIPTOR);
     dwSizeOfFirstThunks += (j+1)*sizeof(IMAGE_THUNK_DATA32)*2; //OriginalFirstThunk + FirstThunk
     
     *pdwSizeOfImportTable = (DWORD)((UCHAR)dwSizeOfDllNames + dwSizeOfFunctionsNames + dwSizeOfFunctionsHints + dwSizeOfFirstThunks + dwSizeOfImportDescriptors);

     //Allocating memory
     pImportTable = (PUCHAR)calloc(*pdwSizeOfImportTable, sizeof(UCHAR));
     
     pdwStartRVA = pImportTableAddress;
     
    
     while(k < i)
     {
         ImportDesc.OriginalFirstThunk = (DWORD)((PUCHAR)pdwStartRVA + dwSizeOfImportDescriptors);
         ImportDesc.TimeDateStamp = 0x00;
         ImportDesc.ForwarderChain = 0x00;
         ImportDesc.Name = (DWORD)((PUCHAR)pdwStartRVA + *pdwSizeOfImportTable - dwSizeOfDllNames);
         ImportDesc.FirstThunk = (DWORD)((PUCHAR)ImportDesc.OriginalFirstThunk + dwSizeOfFirstThunks/2 + sizeof(DWORD));
         
         memcpy((PUCHAR)pImportTable + k * sizeof(IMAGE_IMPORT_DESCRIPTOR), &ImportDesc, sizeof(IMAGE_IMPORT_DESCRIPTOR));
         memcpy((PUCHAR)pImportTable + *pdwSizeOfImportTable - dwSizeOfDllNames, Dlls[0], strlen(Dlls[0]));
         k++;
     }
     
     k = 0;
     
     while(k < j)
     {
         dwFunctionOffset = (PUCHAR)pdwStartRVA + dwSizeOfImportDescriptors + dwSizeOfFirstThunks + 0x2; // 0x2 correspond to Hint : sizeof(WORD)
         if(k > 0)
              (PUCHAR)dwFunctionOffset += strlen(Dlls[k-1])+0x2;
            
         memcpy((PUCHAR)pImportTable + dwSizeOfImportDescriptors + k * sizeof(DWORD), &dwFunctionOffset, sizeof(DWORD));
         memcpy((PUCHAR)pImportTable + 0x2 + dwFunctionOffset - (ULONG)pdwStartRVA, Functions[k], strlen(Functions[k]));
         k++;
     }
     
     k = 0;
     
     while(k < j)
     {
         dwFunctionOffset = (PUCHAR)pdwStartRVA + dwSizeOfImportDescriptors + dwSizeOfFirstThunks + 0x2; // 0x2 correspond to Hint : sizeof(WORD)
         if(k > 0)
              (PUCHAR)dwFunctionOffset += strlen(Dlls[k-1])+0x2;
            
         memcpy((PUCHAR)pImportTable + dwSizeOfImportDescriptors + dwSizeOfFirstThunks/2 +sizeof(DWORD) + k * sizeof(DWORD), &dwFunctionOffset, sizeof(DWORD));
         
         k++;
     }
     
     return pImportTable;
}
VOID DLLEXPORT LPEL_AdjustImportTableAddress(HANDLE hBinary, PDWORD pImportTableAddress)
{
     PIMAGE_IMPORT_DESCRIPTOR pImportTable;
     PIMAGE_THUNK_DATA32 pFirstThunk;
     PDWORD pAddress;
     
     LPEL_SetDirectoryAddress(hBinary, IMAGE_DIRECTORY_ENTRY_IMPORT, pImportTableAddress);
     
     /*
     pImportTable = GetImportTable(hBinary);
     printf("\n\n[+] New Import Table Address : 0x%x", pImportTable);*/

     return;
}
VOID DLLEXPORT LPEL_SetDirectoryAddress(HANDLE hBinary, BYTE bDirectory, PDWORD pAddress)
{
     PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;

     pOptionalHeader = LPEL_GetOptionalHeader(hBinary);

     *(&pOptionalHeader->DataDirectory[bDirectory].VirtualAddress) = pAddress;
      
     return;
}
PDWORD DLLEXPORT LPEL_GetDirectory(HANDLE hBinary, BYTE bDirectory)
{
       PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;

       pOptionalHeader = LPEL_GetOptionalHeader(hBinary);
       
       return (PDWORD)((PUCHAR)LPEL_RVAtoRAW(hBinary, pOptionalHeader->DataDirectory[bDirectory].VirtualAddress) + (ULONG)hBinary);
}
PIMAGE_IMPORT_DESCRIPTOR DLLEXPORT LPEL_GetImportTable(HANDLE hBinary)
{
       return (PIMAGE_IMPORT_DESCRIPTOR)LPEL_GetDirectory(hBinary, IMAGE_DIRECTORY_ENTRY_IMPORT);
} 
DWORD DLLEXPORT LPEL_GetSizeOfDirectory(HANDLE hBinary, BYTE bDirectory)
{
      PIMAGE_OPTIONAL_HEADER pOptionalHeader = NULL;

      pOptionalHeader = LPEL_GetOptionalHeader(hBinary);
       
      return pOptionalHeader->DataDirectory[bDirectory].Size;
}