#include <stdio.h>
#include <fcntl.h>
#include "port.h"
#include "ntbind.h"

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 (ntopt_hdr.CheckSum) {
        puts("can't handle checksum field != 0");
        return (1);
    }

    if (ntopt_hdr.SectionAlignment != 0x10000) {
        puts("this win32 file has not alignment 0x10000");
        return (1);
    } else
        puts("align file from 64K to 4K");

    ntopt_hdr.SectionAlignment = 0x1000;

    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;

    if (argc < 2) {
        puts("align.exe sets win32 section alignment to 4K");
        puts("  usage: align <file>");
        return (1);
    }

    if ((fhandle = open(argv[1], O_RDWR | O_BINARY)) == -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;
}