#include <stdio.h>
#include <sys/mman.h>
#include <string.h>
#include <unistd.h>
#include <elf.h>

#include "log.h"
#include "elf.h"

// Pad ELF File
static void *elf_pad(const char *target, size_t padding, size_t *out_size) {
    // File Data
    FILE *file_obj = NULL;
    unsigned char *file_map = NULL;
    long int file_size = 0;

    // Return
    unsigned char *out = NULL;

    // Code
    {
        // Load Main Binary
        file_obj = fopen(target, "rb");

        // Verify Binary
        if (!file_obj) {
            ERR("Unable To Open: %s", target);
        }

        // Get File Size
        fseek(file_obj, 0L, SEEK_END);
        file_size = ftell(file_obj);
        fseek(file_obj, 0L, SEEK_SET);

        // Map File To Pointer
        file_map = (unsigned char *) mmap(0, file_size, PROT_READ, MAP_PRIVATE, fileno(file_obj), 0);

        // Check ELF Magic
        if (file_map[EI_MAG0] != ELFMAG0 || file_map[EI_MAG1] != ELFMAG1 || file_map[EI_MAG2] != ELFMAG2 || file_map[EI_MAG3] != ELFMAG3) {
            ERR("Not An ELF File: %s", target);
        }
        int is_64bit = file_map[EI_CLASS] != ELFCLASS32;
        if (file_map[EI_DATA] != ELFDATA2LSB) {
            ERR("ELF File Isn't Little-Endian: %s", target);
        }

        // Allocate New File
        *out_size = file_size + padding;
        out = malloc(*out_size);
        ALLOC_CHECK(out);

        // Copy Header
        size_t header_size = is_64bit ? sizeof (Elf64_Ehdr) : sizeof (Elf32_Ehdr);
        memcpy(out, file_map, header_size);
        // Copy Everything Else
        memcpy(out + header_size + padding, file_map + header_size, file_size - header_size);
        // Zero Padding
        memset(out + header_size, 0, padding);

        // Fix Pointers
        if (is_64bit) {
            Elf64_pad(out, padding);
        } else {
            Elf32_pad(out, padding);
        }
   }

   // Unmap And Close File
   if (file_map != NULL) {
       munmap(file_map, file_size);
   }
   if (file_obj != NULL) {
       fclose(file_obj);
   }

   // Return
   return out;
}

// Main
int main(int argc, char *argv[]) {
    if (argc != 3) {
        ERR("%s", "Invalid Arguments");
    }

    // Arguments
    size_t padding = 0;
    if (sscanf(argv[1], "%zu", &padding) != 1) {
        ERR("%s", "Unable To Parse Padding Number");
    }
    char *target = argv[2];

    // Align Padding To Page Boundary
    {
        long page_size = sysconf(_SC_PAGESIZE);
        size_t old_padding = padding;
        padding = (padding / page_size) * page_size;
        if ((old_padding % page_size) != 0) {
            padding += page_size;
        }
    }
    INFO("Aligning Padding To Page Boundary: %zu", padding);

    // Pad
    size_t out_size = 0;
    void *out = elf_pad(target, padding, &out_size);

    // Print
    fwrite(out, 1, out_size, stdout);

    // Free
    free(out);
}