#include <stddef.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>

/*
    int16_t lastsize;
    int16_t nblocks;
    int16_t nreloc;
    int16_t hdrsize;
    int16_t minalloc;
    int16_t maxalloc;
    void *ss; // 2 byte value
    void *sp; // 2 byte value
    int16_t checksum;
    void *ip; // 2 byte value
    void *cs; // 2 byte value
    int16_t relocpos;
    int16_t noverlay;
    int16_t reserved1[4];
    int16_t oem_id;
    int16_t oem_info;
    int16_t reserved2[10];
*/

#define ROUND_UP(N, S) ((((N) + (S) - 1) / (S)) * (S))

struct DOS_Header {
    char signature[2];    // {'M', 'Z'}
    uint16_t  e_cblp;     // Bytes on last page of file
    uint16_t  e_cp;       // Pages in file
    uint16_t  e_crlc;     // Relocations
    uint16_t  e_cparhdr;  // Size of header in paragraphs
    uint16_t  e_minalloc; // Minimum extra paragraphs needed
    uint16_t  e_maxalloc; // Maximum extra paragraphs needed
    uint16_t  e_ss;       // Initial (relative) SS value
    uint16_t  e_sp;       // Initial SP value
    uint16_t  e_csum;     // Checksum
    uint16_t  e_ip;       // Initial IP value
    uint16_t  e_cs;       // Initial (relative) CS value
    uint16_t  e_lfarlc;   // File address of relocation table
    uint16_t  e_ovno;     // Overlay number
    uint16_t  e_res[4];   // Reserved words
    uint16_t  e_oemid;    // OEM identifier (for e_oeminfo)
    uint16_t  e_oeminfo;  // OEM information; e_oemid specific
    uint16_t  e_res2[10]; // Reserved words
    int32_t e_lfanew; // Offset to the 'PE\0\0' signature relative to the beginning of the file
} __attribute__((packed));

struct PE_OpHdr {
    uint16_t format;
    char linker_ver_maj;
    char linker_ver_min;

    uint32_t codesz;
    uint32_t bsssz;
    uint32_t datasz;
    uint32_t entrypoint;
    uint32_t codebase;

    uint64_t imagebase;
    uint32_t sectalign;
    uint32_t filealign;
    uint16_t os_ver_maj;
    uint16_t os_ver_min;
    uint16_t img_ver_maj;
    uint16_t img_ver_min;
    uint16_t subs_ver_maj;
    uint16_t subs_ver_min;
    uint32_t win32_ver;

    uint32_t imagesz;
    uint32_t hdrsz;
    uint32_t cksum;
    uint16_t subsystem;
    uint16_t dllcharacteristics;
    uint64_t stack_res;
    uint64_t stack_commit;
    uint64_t heap_res;
    uint64_t heap_commit;

    uint32_t loader_flags;
    uint32_t nr_rva_sz;
} __attribute__((packed));

struct PE_Header {
    char signature[2];
    uint16_t fill0;
    uint16_t arch;
    uint16_t nr_sections;
    uint32_t datetime;
    uint32_t symtabptr;
    uint32_t nr_symbols;
    uint16_t ophdrsz;
    uint16_t characteristics;
    struct PE_OpHdr ophdr;
} __attribute__((packed));

struct PE_SecHdr {
    char name[8];
    uint32_t virtsz;
    uint32_t vaddr;
    uint32_t rawsz;
    uint32_t dataptr;
    uint32_t relocptr;
    uint32_t linenrptr;
    uint16_t nr_reloc;
    uint16_t nr_linenr;
    uint32_t characteristics;
} __attribute__((packed));

struct PE_Reloc {
    uint32_t offset;
    uint32_t value;
    char fill[0x20-0x8];
} __attribute__((packed));

ssize_t
writeall(int fd, void *buf, size_t sz)
{
    ssize_t wr, ssz;

    ssz = (ssize_t) sz;

    while ((wr = write(fd, buf, sz)) != -1) {
        if (wr == ssz) {
            break;
        }

        buf += wr;
        ssz -= wr;
    }

    return (wr);
}

char *
readall(int fd, size_t *sz)
{
    char *data;
    size_t cap, len;
    ssize_t n;
    int err;

    *sz = 0;

    cap = BUFSIZ;
    len = 0;
    data = malloc(cap);

    if (data == NULL) {
        return (NULL);
    }

    while ((n = read(fd, data+len, BUFSIZ)) >= 0) {
        if (n == 0) {
            break;
        }

        len += n;

        if (len >= cap) {
            cap += BUFSIZ;
            data = realloc(data, cap);
        }
    }
    if (n == -1) {
        err = errno;
        free(data);
        errno = err;
        return (NULL);
    }

    *sz = len;

    return (data);
}

int
main(int argc, char **argv)
{
    struct DOS_Header dhdr;
    struct PE_Header phdr;
    struct PE_SecHdr sect;
    struct PE_Reloc relo;
    uint32_t hdrsz;
    int ifd, ofd, err;
    char *code;
    size_t codesz;
    void *zero;

    assert(offsetof(struct DOS_Header, e_lfanew) == 0x3c);

    if (argc != 3) {
        fprintf(stderr, "Usage: %s INPUT.BIN OUTPUT.BIN\n", argv[0]);
        return (1);
    }

    ifd = open(argv[1], O_RDONLY);
    if (ifd < 0) {
        err = errno;
        perror(argv[1]);
        return (err);
    }

    ofd = open(argv[2], O_WRONLY|O_CREAT|O_TRUNC,
               S_IRUSR|S_IWUSR|S_IRGRP|S_IROTH);

    if (ofd < 0) {
        err = errno;
        perror(argv[2]);
        close(ifd);
        return (err);
    }

    code = readall(ifd, &codesz);
    if (code == NULL) {
        err = errno;
        perror("readall");
        close(ifd);
        close(ofd);
        return (err);
    }

    memset(&dhdr, 0, sizeof(struct DOS_Header));
    memset(&phdr, 0, sizeof(struct PE_Header));
    memset(&sect, 0, sizeof(struct PE_SecHdr));
    memset(&relo, 0, sizeof(struct PE_Reloc));

    dhdr.signature[0] = 'M';
    dhdr.signature[1] = 'Z';
    dhdr.e_lfanew = sizeof(struct DOS_Header);

    writeall(ofd, &dhdr, sizeof(struct DOS_Header));

    phdr.signature[0] = 'P';
    phdr.signature[1] = 'E';
    phdr.arch = 0x8664; /* x86_64 */
    phdr.nr_sections = 2; /* text, reloc */
    phdr.nr_symbols = 1;
    phdr.ophdrsz = sizeof(struct PE_OpHdr);
    printf("phdr = %zu, ophdr = %zu\n", sizeof(struct PE_Header), sizeof(struct PE_OpHdr));
    phdr.characteristics = 0x206;

    phdr.ophdr.format = 0x20b; /* PE32+ (x64) */
    phdr.ophdr.linker_ver_maj = 0x02;
    phdr.ophdr.linker_ver_min = 0x14;
    phdr.ophdr.codebase = 0x200;
    phdr.ophdr.sectalign = 0x20;
    assert(phdr.ophdr.sectalign == sizeof(struct PE_Reloc));
    phdr.ophdr.filealign = 0x20;
    phdr.ophdr.subsystem = 0xa; /* EFI Application */
    phdr.ophdr.nr_rva_sz = 0x0;

    assert(codesz <= (uint32_t)(~0));
    phdr.ophdr.codesz = (uint32_t) codesz;
    printf("codesz = %u\n", phdr.ophdr.codesz);

    hdrsz = sizeof(struct DOS_Header)
        + sizeof(struct PE_Header)
        + phdr.nr_sections * sizeof(struct PE_SecHdr);
    printf("hdrsz = %u\n", hdrsz);
    /* round up to multiple of filealign */
    phdr.ophdr.hdrsz = ROUND_UP(hdrsz, phdr.ophdr.filealign);;
    assert(phdr.ophdr.hdrsz % phdr.ophdr.filealign == 0);
    printf("hdrsz = %u (aligned)\n", phdr.ophdr.hdrsz);

    phdr.ophdr.entrypoint = phdr.ophdr.hdrsz;

    phdr.ophdr.imagesz = phdr.ophdr.hdrsz + codesz + phdr.ophdr.sectalign;
    printf("imagesz = %u\n", phdr.ophdr.imagesz);
    assert(phdr.ophdr.imagesz % phdr.ophdr.sectalign == 0);

    writeall(ofd, &phdr, sizeof(struct PE_Header));

    strcpy((char *)&sect.name, ".reloc");
    /* READ|DISCARDABLE|INITIALIZED */
    sect.characteristics = 0x42000040;
    sect.virtsz = phdr.ophdr.sectalign;
    sect.vaddr = phdr.ophdr.hdrsz + codesz;
    sect.rawsz = phdr.ophdr.sectalign;
    sect.dataptr = phdr.ophdr.hdrsz + codesz;

    relo.offset = sect.dataptr;
    relo.value = 10; /* arbitrary number */

    writeall(ofd, &sect, sizeof(struct PE_SecHdr));

    memset(&sect, 0, sizeof(struct PE_SecHdr));
    strcpy((char *)&sect.name, ".text");
    /* READ|EXEC|ALIGN_16BYTES|CODE */
    sect.characteristics = 0x60500020;
    sect.virtsz = ROUND_UP(codesz, phdr.ophdr.sectalign);
    sect.vaddr = phdr.ophdr.hdrsz;
    sect.rawsz = codesz;
    sect.dataptr = phdr.ophdr.hdrsz;

    writeall(ofd, &sect, sizeof(struct PE_SecHdr));

    if (phdr.ophdr.hdrsz - hdrsz) {
        zero = calloc(1, phdr.ophdr.hdrsz - hdrsz);
        writeall(ofd, zero, phdr.ophdr.hdrsz - hdrsz);
    }

    writeall(ofd, code, codesz);
    writeall(ofd, &relo, sizeof(struct PE_Reloc));

    close(ofd);
    close(ifd);
    free(code);

    return (0);
}