//------------------------------------------------------------------
//  SetEntry.c - Set DLL entry point. Ton Plooy 1998
//------------------------------------------------------------------
#include <stdio.h>
#include <stdlib.h>
#define  WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <winnt.h>

//------------------------------------------------------------------
LPVOID MapFile(char *szFileName, HANDLE *hFile, HANDLE *hMap);
void   UnmapFile(HANDLE hFile, HANDLE hMap, LPVOID pMem);
DWORD  ProcessDLL(LPBYTE pIn, char *szFunction);
DWORD  GetExportRVA(LPBYTE pSect, DWORD SectRva, char *szFunction);
DWORD  FindRVA(LPBYTE pFileMap, IMAGE_NT_HEADERS *pHdr, DWORD Rva);

//------------------------------------------------------------------

int main(int argc, char *argv[])
{
    // Set the entry point of the specified PE binary to the
    // given function. The function must have an export.
    HANDLE hFile, hMap;
    LPVOID pIn;
    DWORD  Stat;

    if (argc != 3) {
        printf("Usage: SetEntry <DllName> <DllEntryPoint>\n");
        return ERROR_INVALID_PARAMETER;
    }
    printf("SetEntry v1.00\n");

    // Make a (memory) file mapping for the binary
    if ((pIn = MapFile(argv[1], &hFile, &hMap)) == NULL) {
        printf("Error: couldn't open or map %s\n", argv[1]);
        return ERROR_OPEN_FAILED;
    }

    if ((Stat = ProcessDLL((LPBYTE)pIn, argv[2])) != NO_ERROR) {
        printf("Error %d in processing %s\n", Stat, argv[1]);
    }

    // Unmap it, memory changes will be written to the image
    UnmapFile(hFile, hMap, pIn);

    return NO_ERROR;
}

//------------------------------------------------------------------

LPVOID MapFile(char *szFile, HANDLE *hFile, HANDLE *hMap)
{
    // Create a memory mapping for the specified file
    DWORD Size;

    *hFile = *hMap = NULL;
    if ((*hFile = CreateFile(szFile, GENERIC_READ | GENERIC_WRITE,
                    FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
                    OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)) ==
       INVALID_HANDLE_VALUE) {
        return NULL;
    }

    Size = GetFileSize(*hFile, NULL);

    if ((*hMap = CreateFileMapping(*hFile, NULL, PAGE_READWRITE, 0,
       Size, NULL)) == NULL) {
        CloseHandle(hFile);
        return NULL;
    }
    return MapViewOfFile(*hMap, FILE_MAP_ALL_ACCESS, 0, 0, 0);
}

//------------------------------------------------------------------

void UnmapFile(HANDLE hFile, HANDLE hMap, LPVOID pMem)
{
    if (pMem != NULL)
        UnmapViewOfFile(pMem);
    if (hMap != NULL)
        CloseHandle(hMap);
    if (hFile != NULL)
        CloseHandle(hFile);
    return ;
}

//------------------------------------------------------------------

DWORD ProcessDLL(LPBYTE pIn, char *szFunction)
{
    long                  lPeStart;
    DWORD                 ExpRVA, ExpOff, FuncAddr;
    IMAGE_NT_HEADERS     *pNtHdr;
    IMAGE_DATA_DIRECTORY *pDataDir;

    // See where the PE image header is located
    if ((((IMAGE_DOS_HEADER *)pIn)->e_magic
   != IMAGE_DOS_SIGNATURE) ||
        ((lPeStart = ((IMAGE_DOS_HEADER *)pIn)->e_lfanew) == 0L)) {
        return ERROR_BAD_EXE_FORMAT;
    };

    pNtHdr = (IMAGE_NT_HEADERS *)(pIn + lPeStart);
    if (pNtHdr->Signature != IMAGE_NT_SIGNATURE) {
        return ERROR_INVALID_EXE_SIGNATURE;
    }

    // Read RVA of the Exports Section from the IMAGE_DIRECTORY
    pDataDir = pNtHdr->OptionalHeader.DataDirectory;
    ExpRVA = pDataDir[IMAGE_DIRECTORY_ENTRY_EXPORT].VirtualAddress;

    // Find file offset for export section RVA
    if ((ExpOff = FindRVA(pIn, pNtHdr, ExpRVA)) == 0) {
        printf("Export section not found\n");
        return ERROR_PROC_NOT_FOUND;
    }
    // Get function address from export section
    FuncAddr = GetExportRVA(pIn + ExpOff, ExpRVA, szFunction);
    if (FuncAddr == 0) {
        printf("Unable to locate export for %s\n", szFunction);
        return ERROR_PROC_NOT_FOUND;
    }

    // Now set the new entry point in the PE header
    printf("Old Entry Point: %08lX\n",
        pNtHdr->OptionalHeader.AddressOfEntryPoint);

    pNtHdr->OptionalHeader.AddressOfEntryPoint = FuncAddr;

    printf("New Entry Point: %08lX\n",
        pNtHdr->OptionalHeader.AddressOfEntryPoint);

    return NO_ERROR;
}

//------------------------------------------------------------------
    
DWORD GetExportRVA(LPBYTE pExpSect, DWORD SectRva, char *szFunction)
{
    // Scan the export section to see if the specified function name
    // can be found. If found, we return its RVA.
    DWORD                    i = 0;
    IMAGE_EXPORT_DIRECTORY  *pDir;
    DWORD                   *pNames, *pAddr;
    WORD                    *pOrd;

    pDir   = (IMAGE_EXPORT_DIRECTORY *)pExpSect;
    pNames = (DWORD *)(pExpSect +
                        (DWORD)pDir->AddressOfNames - SectRva);
    printf("Searching exports..\n");
    while ((i < pDir->NumberOfNames) && (strcmp(pExpSect +
                        (pNames[i] - SectRva), szFunction) != 0)) {
        printf("%s\n", pExpSect + (pNames[i] - SectRva));
        i++;
    }
    if (i == pDir->NumberOfNames)
        return 0; // Function not found

    // Get Function ordinal
    pOrd = (WORD *)(pExpSect +
                    (DWORD)pDir->AddressOfNameOrdinals - SectRva);

    // Get function address table. Function RVA is at ordinal index
    pAddr = (DWORD *)(pExpSect +
                    (DWORD)pDir->AddressOfFunctions - SectRva);

    return pAddr[pOrd[i]];
}

//------------------------------------------------------------------

DWORD FindRVA(LPBYTE pFileMap, IMAGE_NT_HEADERS *pNtHdr, DWORD Rva)
{
    // Find the file byte offset for the given RVA
    int i         = 0,
        nSections = (int)pNtHdr->FileHeader.NumberOfSections;
    IMAGE_SECTION_HEADER *pSectHdr;

    pSectHdr = (IMAGE_SECTION_HEADER *)((LPSTR)pNtHdr +
   sizeof(IMAGE_NT_HEADERS));
    while ((i < nSections) && ((pSectHdr[i].VirtualAddress +
                    pSectHdr[i].Misc.VirtualSize - 1) < Rva)) {
        i++;
    }
    if (((i == nSections) && ((pSectHdr[i].VirtualAddress +
                    pSectHdr[i].Misc.VirtualSize - 1) < Rva)) ||
                    (pSectHdr[i].VirtualAddress > Rva)) {
        return 0;
    }
    return pSectHdr[i].PointerToRawData +
                            Rva - pSectHdr[i].VirtualAddress;
}
// End of File