/**
 * BNCSutil
 * Battle.Net Utility Library
 *
 * Copyright (C) 2004-2006 Eric Naeseth
 *
 * Portable Executable Processing
 * August 16, 2005
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * A copy of the GNU Lesser General Public License is included in the BNCSutil
 * distribution in the file COPYING.  If you did not receive this copy,
 * write to the Free Software Foundation, Inc., 59 Temple Place, Suite 330,
 * Boston, MA  02111-1307  USA
 */

#include <bncsutil/pe.h>
#include <bncsutil/stack.h>
#include <stdlib.h>
#include <string.h>

int cm_pe_load_resdir(FILE* f, uint32_t offset, cm_pe_resdir_t* dir);

MEXP(cm_pe_t) cm_pe_load(const char* filename)
{
    cm_pe_t pe;
    cm_pe_header_t* header;
    cm_pe_optional_header_t* opt_header;
    cm_pe_windows_header_t* win_header;
    long pe_offset = 0;
    //size_t i;
    size_t count;

    pe = (cm_pe_t) malloc(sizeof(struct cm_pe));
    if (!pe)
        return (cm_pe_t) 0;

    memset(pe, 0, sizeof(struct cm_pe));

    pe->f = fopen(filename, "r");
    if (!pe->f) {
        free(pe);
        return (cm_pe_t) 0;
    }

    if (fseek(pe->f, 0x3CL, SEEK_SET) == -1)
        goto err_trap;

    if (fread(&pe_offset, 4, 1, pe->f) != 1)
        goto err_trap;

#if BIG_ENDIAN
    pe_offset = LSB4(pe_offset);
#endif

    if (fseek(pe->f, pe_offset, SEEK_SET) == -1)
        goto err_trap;

    if (fread(&pe->header, sizeof(cm_pe_header_t), 1, pe->f) != 1)
        goto err_trap;

    header = &pe->header;
#if BIGENDIAN
    /* Without regular expressions, this would've sucked. */
    header->signature = SWAP32(header->signature);
    header->machine = SWAP16(header->machine);
    header->section_count = SWAP16(header->section_count);
    header->timestamp = SWAP32(header->timestamp);
    header->symbol_table_offset = SWAP32(header->symbol_table_offset);
    header->symbol_count = SWAP32(header->symbol_count);
    header->optional_header_size = SWAP16(header->optional_header_size);
    header->characteristics = SWAP16(header->characteristics);
#endif

    if (header->optional_header_size > 0) {
        if (fread(&pe->optional_header, PE_OPTIONAL_HEADER_MIN_SIZE, 1, pe->f)
            != 1)
        {
            goto err_trap;
        }

        opt_header = &pe->optional_header;
        win_header = &pe->windows_header;

#if BIGENDIAN
        opt_header->magic = SWAP16(opt_header->magic);
#endif

        if (opt_header->magic == IMAGE_FORMAT_PE32) {
            if (fread(&opt_header->data_base, 4, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->image_base, 4, 1, pe->f) != 1)
                goto err_trap;
            // The 40 is not a typo.
            if (fread(&win_header->section_alignment, 40, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->stack_reserve_size, 4, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->stack_commit_size, 4, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->heap_reserve_size, 4, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->heap_commit_size, 4, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->loader_flags, 4, 1, pe->f) != 1)
                goto err_trap;
            if (fread(&win_header->data_directory_count, 4, 1, pe->f) != 1)
                goto err_trap;
        } else if (opt_header->magic == IMAGE_FORMAT_PE32_PLUS) {
            if (fread(win_header, sizeof(cm_pe_windows_header_t), 1, pe->f)!= 1)
                goto err_trap;
        } else {
            goto err_trap;
        }

#if BIGENDIAN
        opt_header->code_section_size = SWAP32(opt_header->code_section_size);
        opt_header->initialized_data_size =
            SWAP32(opt_header->initialized_data_size);
        opt_header->uninitialized_data_size =
            SWAP32(opt_header->uninitialized_data_size);
        opt_header->entry_point = SWAP32(opt_header->entry_point);
        opt_header->code_base = SWAP32(opt_header->code_base);
        opt_header->data_base = SWAP32(opt_header->data_base);

        win_header->image_base = SWAP64(win_header->image_base);
        win_header->section_alignment = SWAP32(win_header->section_alignment);
        win_header->file_alignment = SWAP32(win_header->file_alignment);
        win_header->major_os_version = SWAP16(win_header->major_os_version);
        win_header->minor_os_version = SWAP16(win_header->minor_os_version);
        win_header->major_image_version =
            SWAP16(win_header->major_image_version);
        win_header->minor_image_version =
            SWAP16(win_header->minor_image_version);
        win_header->major_subsystem_version =
            SWAP16(win_header->major_subsystem_version);
        win_header->minor_subsystem_version =
            SWAP16(win_header->minor_subsystem_version);
        win_header->reserved = SWAP32(win_header->reserved);
        win_header->image_size = SWAP32(win_header->image_size);
        win_header->headers_size = SWAP32(win_header->headers_size);
        win_header->checksum = SWAP32(win_header->checksum);
        win_header->subsystem = SWAP16(win_header->subsystem);
        win_header->dll_characteristics =
            SWAP16(win_header->dll_characteristics);
        win_header->stack_reserve_size = SWAP64(win_header->stack_reserve_size);
        win_header->stack_commit_size = SWAP64(win_header->stack_commit_size);
        win_header->heap_reserve_size = SWAP64(win_header->heap_reserve_size);
        win_header->heap_commit_size = SWAP64(win_header->heap_commit_size);
        win_header->loader_flags = SWAP32(win_header->loader_flags);
        win_header->data_directory_count =
            SWAP32(win_header->data_directory_count);
#endif

        if (win_header->data_directory_count > 0) {
            count = win_header->data_directory_count;
            pe->data_directories = (cm_pe_data_directory_t*)
                calloc(sizeof(cm_pe_data_directory_t), count);

            if (!pe->data_directories)
                goto err_trap;

            if (fread(pe->data_directories, sizeof(cm_pe_data_directory_t),
                count, pe->f) != count)
            {
                goto dir_err_trap;
            }

#if BIGENDIAN
            for (i = 0; i < count; i++) {
                pe->data_directories[i].rva =
                    SWAP32(pe->data_directories[i].rva);
                pe->data_directories[i].size =
                    SWAP32(pe->data_directories[i].size);
            }
#endif
        }

        count = (size_t) header->section_count;
        if (count) {
            pe->sections = (cm_pe_section_t*) calloc(sizeof(cm_pe_section_t),
                                                     count);

            if (!pe->sections)
                goto dir_err_trap;

            if (fread(pe->sections, sizeof(cm_pe_section_t), count, pe->f)
                != count)
            {
                goto sect_err_trap;
            }
        }

#if BIGENDIAN
        for (i = 0; i < count; i++) {
            pe->sections[i].virtual_size = SWAP32(pe->sections[i].virtual_size);
            pe->sections[i].virtual_address =
                SWAP32(pe->sections[i].virtual_address);
            pe->sections[i].raw_data_size =
                SWAP32(pe->sections[i].raw_data_size);
            pe->sections[i].raw_data_offset =
                SWAP32(pe->sections[i].raw_data_offset);
            pe->sections[i].relocations_offset =
                SWAP32(pe->sections[i].relocations_offset);
            pe->sections[i].line_numbers_offset =
                SWAP32(pe->sections[i].line_numbers_offset);
            pe->sections[i].relocation_count =
                SWAP16(pe->sections[i].relocation_count);
            pe->sections[i].line_number_count =
                SWAP16(pe->sections[i].line_number_count);
            pe->sections[i].characteristics =
                SWAP32(pe->sections[i].characteristics);
        }
#endif
    }

    return pe;
sect_err_trap:
    free(pe->sections);
dir_err_trap:
    free(pe->data_directories);
err_trap:
    fclose(pe->f);
    free(pe);
    return (cm_pe_t) 0;
}

MEXP(void) cm_pe_unload(cm_pe_t pe)
{
    if (pe->data_directories)
        free(pe->data_directories);
    if (pe->sections)
        free(pe->sections);
    if (pe->f)
        fclose(pe->f);
    free(pe);
}

static cm_pe_section_t* cm_pe_get_section(cm_pe_t pe, const char* name) {
    unsigned int i;
    cm_pe_section_t* s;
    uint32_t section_count = pe->header.section_count;

    if (!pe || !pe->sections)
        return (cm_pe_section_t*) 0;

    for (i = 0, s = pe->sections; i < section_count; i++, s++) {
        if (strcmp(s->name, name) == 0)
            return s;
    }

    return (cm_pe_section_t*) 0;
}

MEXP(cm_pe_resdir_t*) cm_pe_load_resources(cm_pe_t pe)
{
    cm_pe_section_t* sect;
    cm_pe_resdir_t* root = (cm_pe_resdir_t*) 0;
    cm_pe_resdir_t* dir;
    cm_pe_resdir_t* subdirs;
    cm_pe_res_t res;
    cm_pe_res_t* resources;
    cm_stack_t stack;
    size_t i;
    uint32_t base;

    // no need to check validity of pe pointer; cm_pe_get_section does this
    sect = cm_pe_get_section(pe, ".rsrc");
    if (!sect)
        return (cm_pe_resdir_t*) 0;

    root = (cm_pe_resdir_t*) malloc(sizeof(cm_pe_resdir_t));
    if (!root)
        return (cm_pe_resdir_t*) 0;

    base = sect->raw_data_offset;
    if (!cm_pe_load_resdir(pe->f, base, root)) {
        free(root);
        return (cm_pe_resdir_t*) 0;
    }

    stack = cm_stack_create();
    if (!stack) {
        free(root);
        return (cm_pe_resdir_t*) 0;
    }

    cm_stack_push(stack, root);

    while ( (dir = (cm_pe_resdir_t*) cm_stack_pop(stack)) ) {
        while (dir->subdir_count + dir->resource_count <
               dir->named_entry_count + dir->id_entry_count)
        {
            if (fseek(pe->f, dir->offset, SEEK_SET) == -1) {
                cm_pe_unload_resources(root);
                cm_stack_destroy(stack);
                return (cm_pe_resdir_t*) 0;
            }

            if (fread(&res, CM_RES_REAL_SIZE, 1, pe->f) != 1) {
                cm_pe_unload_resources(root);
                cm_stack_destroy(stack);
                return (cm_pe_resdir_t*) 0;
            }

    #if BIGENDIAN
            res.name = SWAP4(res.name);
            res.offset = SWAP4(res.offset);
    #endif
            if (res.offset & 0x80000000) {
                // subdirectory
                i = dir->subdir_count++;
                subdirs = (cm_pe_resdir_t*) realloc(dir->subdirs,
                    sizeof(cm_pe_resdir_t) * dir->subdir_count);
                if (!subdirs) {
                    cm_pe_unload_resources(root);
                    cm_stack_destroy(stack);
                    return (cm_pe_resdir_t*) 0;
                }
                dir->subdirs = subdirs;

                cm_stack_push(stack, dir);
                dir->offset += CM_RES_REAL_SIZE;
                dir = (subdirs + i);

                res.offset &= 0x7FFFFFFF;
                res.file_offset = base + res.offset;

                if (!cm_pe_load_resdir(pe->f, res.file_offset, dir)) {
                    cm_pe_unload_resources(root);
                    cm_stack_destroy(stack);
                    return (cm_pe_resdir_t*) 0;
                }
                dir->name = res.name;

                cm_stack_push(stack, dir);
                break;
            }
            // real resource
            res.file_offset = base + res.offset;
            i = dir->resource_count++;
            resources = (cm_pe_res_t*) realloc(dir->resources,
                sizeof(cm_pe_res_t) * dir->resource_count);
            if (!resources) {
                cm_pe_unload_resources(root);
                cm_stack_destroy(stack);
                return (cm_pe_resdir_t*) 0;
            }
            dir->resources = resources;
            memcpy(dir->resources + i, &res, sizeof(cm_pe_res_t));
            dir->offset += CM_RES_REAL_SIZE;
        }
    }

    cm_stack_destroy(stack);
    return root;
}

MEXP(int) cm_pe_unload_resources(cm_pe_resdir_t* root)
{
    cm_pe_resdir_t* dir;
    cm_stack_t stack;

    stack = cm_stack_create();
    if (!stack)
        return 0;

    cm_stack_push(stack, root);

    while ( (dir = cm_stack_pop(stack)) ) {
        if (dir->subdir_count) {
            dir->subdir_count--;
            cm_stack_push(stack, dir);
            cm_stack_push(stack, (dir->subdirs + dir->subdir_count));
            continue;
        }

        if (dir->subdirs) {
            free(dir->subdirs);
            dir->subdirs = (cm_pe_resdir_t*) 0;
        }

        if (dir->resources) {
            free(dir->resources);
            dir->resource_count = 0;
        }
    }

    cm_stack_destroy(stack);
    free(root);
    return 1;
}

MEXP(int) cm_pe_fixed_version(cm_pe_t pe, cm_pe_res_t* res,
                              cm_pe_version_t* ver)
{
    // find ".rsrc" section
    cm_pe_section_t* sect = pe->sections;
    {
        int i=0;
        for(; i<pe->header.section_count; ++i)
        {
            if(strncmp(sect[i].name,".rsrc",sizeof(sect[i].name))==0)
            {
                sect = &sect[i];
                break;
            }
        }
        if(i == pe->header.section_count) // no resources ???
            return 0;
    }
#if BIGENDIAN
    uint32_t check = 0xBD04EFFE;
#else
    uint32_t check = 0xFEEF04BD;
#endif
    uint32_t rva;
    uint32_t size;
    uint32_t offset;

    if (!pe || !res || !ver)
        return 0;

    if (fseek(pe->f, res->file_offset, SEEK_SET) == -1)
        return 0;
    if (fread(&rva, 4, 1, pe->f) != 1)
        return 0;
    if (fread(&size, 4, 1, pe->f) != 1)
        return 0;
#if BIGENDIAN
    rva = SWAP4(rva);
    size = SWAP4(size);
#endif

    offset = sect->raw_data_offset + (rva - sect->virtual_address) + 0x26;
    offset = (offset+3) & 0xFFFFFFFC; // align 4 byte
    if (fseek(pe->f, offset, SEEK_SET) == -1)
        return 0;
    if (fread(ver, sizeof(cm_pe_version_t), 1, pe->f) != 1)
        return 0;

    if (ver->dwSignature != check)
        return 0;

#if BIGENDIAN
    ver->dwSignature = SWAP32(ver->dwSignature);
    ver->dwStrucVersion = SWAP32(ver->dwStrucVersion);
    ver->dwFileVersionMS = SWAP32(ver->dwFileVersionMS);
    ver->dwFileVersionLS = SWAP32(ver->dwFileVersionLS);
    ver->dwProductVersionMS = SWAP32(ver->dwProductVersionMS);
    ver->dwProductVersionLS = SWAP32(ver->dwProductVersionLS);
    ver->dwFileFlagsMask = SWAP32(ver->dwFileFlagsMask);
    ver->dwFileFlags = SWAP32(ver->dwFileFlags);
    ver->dwFileOS = SWAP32(ver->dwFileOS);
    ver->dwFileType = SWAP32(ver->dwFileType);
    ver->dwFileSubtype = SWAP32(ver->dwFileSubtype);
    ver->dwFileDateMS = SWAP32(ver->dwFileDateMS);
    ver->dwFileDateLS = SWAP32(ver->dwFileDateLS);
#endif

    return 1;
}

int cm_pe_load_resdir(FILE* f, uint32_t offset, cm_pe_resdir_t* dir)
{
    memset(dir, 0, sizeof(cm_pe_resdir_t));

    if (fseek(f, offset, SEEK_SET) == -1)
        return 0;

    if (fread(dir, 16, 1, f) != 1)
        return 0;

#if BIGENDIAN
    dir->characteristics = SWAP32(dir->characteristics);
    dir->timestamp = SWAP32(dir->timestamp);
    dir->major_version = SWAP16(dir->major_version);
    dir->minor_version = SWAP16(dir->minor_version);
    dir->named_entry_count = SWAP16(dir->named_entry_count);
    dir->id_entry_count = SWAP16(dir->id_entry_count);
#endif

    dir->offset = (uint32_t) ftell(f);

    return 1;
}
