/*

    RELOC.C - relocates Win32 executable to new base address

    Copyright (C) 1995-1996
	Rainer Schnitker, Heeper Str. 283, 33607 Bielefeld
	email: rainer@mathematik.uni-bielefeld.de

    All rights reserved

*/

#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <malloc.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include "port.h"
#include "ntbind.h"

#define ILLEGAL_BASE 0xffffffff

DWORD text_offset, text_addr, text_end;
DWORD data_offset, data_addr, data_end;
DWORD reloc_offset, reloc_addr, reloc_size;

static int reloc_image(int fhandle, long dx)
{
    struct _IMAGE_BASE_RELOCATION {
	DWORD	VirtualAddress;
	DWORD	SizeOfBlock;
    } *baseReloc, *save;
    char *image;

    image = (char *) malloc(4096+4);
    if (!image)
	return -1;

    save = baseReloc = (struct _IMAGE_BASE_RELOCATION *) malloc(reloc_size);
    if (!baseReloc)
	return -1;

    lseek(fhandle, reloc_offset, 0);
    read(fhandle, baseReloc, reloc_size);

    for (;;) {
	unsigned i,cEntries;
	WORD relocType;
	WORD *pEntry;
	long newpos;

        if (!baseReloc->VirtualAddress || !baseReloc->SizeOfBlock)
            break;

        printf("Vma=%08lX entries=%lX\n",
                baseReloc->VirtualAddress, baseReloc->SizeOfBlock);

        if (baseReloc->VirtualAddress >= data_addr &&
            baseReloc->VirtualAddress <= data_end)
	    newpos = data_offset + baseReloc->VirtualAddress - data_addr;
        else if (baseReloc->VirtualAddress >= text_addr &&
                 baseReloc->VirtualAddress <= text_end)
	    newpos = text_offset + baseReloc->VirtualAddress - text_addr;
        else
            newpos = 0;  /* bss */

        if (newpos) {
            lseek (fhandle, newpos, 0);
            read (fhandle, image, 4096 + 4);

            cEntries = (baseReloc->SizeOfBlock -
                        sizeof(struct _IMAGE_BASE_RELOCATION)) / sizeof(WORD);
            pEntry = (WORD *) (baseReloc + 1);

            for ( i=0; i < cEntries; i++ )
            {
                relocType = (*pEntry & 0xF000) >> 12;

                if (relocType) {
                    DWORD offset = (*pEntry & 0x0FFF);
                    DWORD value = *(DWORD *) (image + offset);

                    *(DWORD *)(image + offset) = value + dx;
                }

                pEntry++;
            }
            lseek (fhandle, newpos, 0);
            write (fhandle, image, 4096 + 4);
        }
        baseReloc = (struct _IMAGE_BASE_RELOCATION *)
		((char *) baseReloc + baseReloc->SizeOfBlock);
    }

    free(image);
    free(save);
    return 0;
}

static int find_pe_signatur(char *program, int filehandle)
{
    struct exe_hdr exehdr;
    struct emx_hdr emxhdr;
    DWORD new_off, signatur;

    lseek(filehandle, 0, 0);
    read(filehandle, &exehdr, sizeof(struct exe_hdr));

    if (exehdr.signatur != 0x5a4d)
	return 0;

    lseek(filehandle, ((long)exehdr.hdr_para) * 16, SEEK_SET);
    read(filehandle, &emxhdr, sizeof(emxhdr));

    if (memcmp(emxhdr.sig, "emx", 3) == 0) {
	printf("%s: cannot handle emx/dual file type\n", program);
	return 0;
    }

    lseek(filehandle, 0x3c, 0);
    read(filehandle, &new_off, 4);
    lseek(filehandle, new_off, 0);
    read(filehandle, &signatur, 4);

    if (signatur != 0x00004550) {
	printf("%s is not a win32 file (PE)\n", program);
	return 0;
    }

    return 1;
}

static int reloc_sections(char *program, DWORD newbase)
{
    FILEHDR	    file_hdr;
    NTOPTHDR	    nt_hdr;
    SCNHDR	    scn_hdr;
    int 	    pehandle;
    int 	    i;
    long	    nt_hdr_pos;
    int             mode;

    if (newbase == ILLEGAL_BASE)
        mode = O_RDONLY;
    else
        mode = O_RDWR;

    if ((pehandle = open(program, mode | O_BINARY)) == -1) {
	perror(program);
	return -1;
    }

    if (!find_pe_signatur(program, pehandle)) {
	close(pehandle);
	return -1;
    }

    read(pehandle, &file_hdr, sizeof(file_hdr));
    if (file_hdr.f_magic != 0x14C) {	/* COFF header */
	printf("%s is not a win32 file (missing coff)\n", program);
	close(pehandle);
	return -1;
    }

    nt_hdr_pos = tell(pehandle);
    read(pehandle, &nt_hdr, sizeof(NTOPTHDR));

    if (nt_hdr.ImageBase == newbase || newbase == ILLEGAL_BASE) {
        printf("%s base address is %s %lX\n",
                program,
                (newbase != ILLEGAL_BASE) ? "already:" : ":",
                nt_hdr.ImageBase);
	close(pehandle);
	return 0;
    }

    text_offset = data_offset = reloc_offset = 0;

    for (i = 0; i < file_hdr.f_nscns; i++) {
	read(pehandle, &scn_hdr, sizeof(SCNHDR));

	if (strcmp(scn_hdr.s_name, ".text") == 0) {
	    text_offset = scn_hdr.s_scnptr;
	    text_addr = scn_hdr.s_vaddr;
            text_end = text_addr + scn_hdr.s_size;
	}
	else if (strcmp(scn_hdr.s_name, ".data") == 0) {
	    data_offset = scn_hdr.s_scnptr;
	    data_addr = scn_hdr.s_vaddr;
            data_end = data_addr + scn_hdr.s_size;
	}
	else if (strcmp(scn_hdr.s_name, ".reloc") == 0) {
	    reloc_offset = scn_hdr.s_scnptr;
	    reloc_addr = scn_hdr.s_vaddr;
	    reloc_size = scn_hdr.s_size;
	}
        else if (strcmp(scn_hdr.s_name, ".rsrc") == 0) {
            if (newbase != ILLEGAL_BASE) {
                printf("RELOC: this tool cannot relocate resources!\n");
                printf("You should first relocate the program image base\n");
                printf("and then you can add the resources with rsrc\n");
                return -1;
            }
        }
    }

    if (!text_offset || !data_offset || !reloc_offset) {
        puts("Bad sections: missing .text or .data or .reloc");
	return -1;
    }

    if (reloc_image(pehandle, newbase - nt_hdr.ImageBase) < 0) {
        static char msg[] = "out of memory";

        write(1, msg, sizeof(msg) - 1);
	close(pehandle);
	return -1;
    }

    /* write new NT image base */
    nt_hdr.ImageBase = newbase;
    lseek(pehandle, nt_hdr_pos, 0);
    write(pehandle, &nt_hdr, sizeof(NTOPTHDR));
    close(pehandle);

    printf("%s ok\n", program);
    return 0;
}

long hex_atol(char *s)
{
    int i;
    long res = 0;

    for (i = 0; i < 8; i++) {
	char c = s[i];

	if (!c)
	    break;
	if (c >= 'a')
	    c -= ('a' - 10);
	else if (c >= 'A')
	    c -= ('A' - 10);
	else
	    c -= '0';
	res <<= 4;
	res |= c;
    }
    return res;
}

int main(int argc, char **argv)
{
    static char usage_text[] =
	"usage:\n"
        " reloc [-Base=<base>] win32-files (wildcards)\n"
	"examples:\n"
	" reloc -Base=0x401000 test.dll\n"
	" reloc *.exe\n";

    DWORD base;
    int i;
    char base_asc[20];

#if defined (__EMX__)
    _envargs (&argc, &argv, "ARGS");
    _response (&argc, &argv);
    _wildcard (&argc, &argv);
#endif

    if (argc == 1) {
	puts(usage_text);
	return -1;
    }

    base = ILLEGAL_BASE;

    for (i = 1; i < argc; i++)
	if (argv[i][0] == '-') {
	    if (strncmp(argv[i], "-Base=", 6) == 0) {
		if (strlen(argv[i] + 6) <= 10) {
		    strcpy(base_asc, argv[i] + 6);
		    strlwr(base_asc);
		    if (strncmp(base_asc, "0x", 2) == 0)
			base = hex_atol(base_asc + 2);
		    else
			base = atol(base_asc);
		}
		else {
		    puts("illegal base value");
		    return 1;
		}
	    }
	    else {
		puts("illegal option");
		puts(usage_text);
		return 1;
	    }
	}

    for (i = 1; i < argc; i++)
	if (argv[i][0] != '-')
	    reloc_sections(argv[i], base);

    return 0;
}
