/*

    PESTACK.C - modifies stack size

    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 <fcntl.h>
#include "port.h"
#include "ntbind.h"

unsigned long stack, commit;

static int modify_header(int pehandle, DWORD headoff)
{
    static FILEHDR         file_hdr;
    static NTOPTHDR        ntopt_hdr;

    lseek(pehandle, headoff, SEEK_SET);
    read(pehandle, &file_hdr, sizeof(FILEHDR));

    if (file_hdr.f_magic != 0x14C) {
        puts("no COFF header found");
        return (1);
    }

    read(pehandle, &ntopt_hdr, sizeof(NTOPTHDR));

    if (ntopt_hdr.Magic != 0x10b) {
        puts("no optional header found");
        return (1);
    }

    if (!stack && !commit) {    /* only print */
        printf("StackReserve 0x%lX = %ld K\nStackCommit 0x%lX = %ld K",
            ntopt_hdr.SizeOfStackReserve,
            ntopt_hdr.SizeOfStackReserve / 1024,
            ntopt_hdr.SizeOfStackCommit,
            ntopt_hdr.SizeOfStackCommit / 1024);
        return 0;
    }

    if (ntopt_hdr.CheckSum) {
        puts("can't handle checksum field != 0");
        return (1);
    }

    ntopt_hdr.SizeOfStackReserve = stack;
    ntopt_hdr.SizeOfStackCommit = commit;

    lseek(pehandle, headoff + sizeof(FILEHDR), SEEK_SET);
    write(pehandle, &ntopt_hdr, sizeof(NTOPTHDR));

    return 0;
}

static int skip_exe_hdr(int filehandle, DWORD * headoff)
{
    struct exe_hdr exehdr;

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

    if (exehdr.signatur == 0x5a4d) {	/* falls exe-kopf */
        DWORD new_off;

        lseek(filehandle, 0x3C, SEEK_SET);
        read(filehandle, &new_off, sizeof(long));

        if (new_off) {
            DWORD pe_magic;
            lseek(filehandle, new_off, SEEK_SET);
            read(filehandle, &pe_magic, sizeof(pe_magic));
            if (pe_magic == 0x4550) {
                *headoff = new_off + 4;
                return 0;
            }
        }
    }   /* exe files */
    *headoff = 0;
    return -1;
}

int main(int argc, char **argv)
{
    int fhandle;
    DWORD headoff;
    int i;
    int mode;

    if (argc == 2) {
        i = 1;
        mode = O_RDONLY | O_BINARY;
    }
    else if (argc == 4 && !memcmp(argv[1], "-s", 2) && !memcmp(argv[2], "-c", 2)) {
        i = 3;
        mode = O_RDWR | O_BINARY;
        stack = atol(argv[1] + 2) * 1024L;
        commit = atol(argv[2] + 2) * 1024L;

        stack = (stack / 4) * 4;
        commit = (commit / 4) * 4;
    }
    else  {
        puts("pestack.exe sets/view the stack size of Win32 apps");
        puts("  usage: pestack [-s<size in K> -c<commit in K>] <file>");
        return (1);
    }

    if ((fhandle = open(argv[i], mode)) == -1) {
        perror("open");
        return (1);
    }

    if (skip_exe_hdr(fhandle, &headoff) < 0) {
        puts("no a valid Win32 file");
        return (1);
    }

    modify_header(fhandle, headoff);

    close(fhandle);

    return 0;
}
