#include <inttypes.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <err.h>
#include <string.h>
#include <mach-o/arch.h>
#include <mach-o/loader.h>
#include <mach-o/nlist.h>

typedef struct macho_input {
    const void *data;
    size_t length;
} macho_input_t;

struct _dyld {
    const struct dyld_info_command *dyld_cmd;
    char *name;
};

/// Verify that the given range is within bounds
static const void *macho_read (macho_input_t *input, const void *address, size_t length) {
    if ((((uint8_t *) address) - ((uint8_t *) input->data)) + length > input->length) {
        warnx("Short read parsing Mach-O input");
        return NULL;
    }
    return address;
}

// Verify that address + offset + length is within bounds.
static const void *macho_offset (macho_input_t *input, const void *address, size_t offset, size_t length) {
    void *result = ((uint8_t *) address) + offset;
    return macho_read(input, result, length);
}

uintptr_t get_dyld_count(macho_input_t *input, const struct mach_header_64 *header, const struct load_command *cmd) {
    uintptr_t ndyld = 0;
    for (uintptr_t i = 0; i < header->ncmds; i++) {
        uint32_t cmdsize = cmd->cmdsize;
        cmd = macho_read(input, cmd, cmdsize);

        switch(cmd->cmd) {
            case LC_ID_DYLIB:
            case LC_LOAD_WEAK_DYLIB:
            case LC_REEXPORT_DYLIB:
            case LC_LOAD_DYLIB: {
                ndyld++;
            }
        }

        cmd = macho_offset(input, cmd, cmdsize, sizeof(struct load_command));
        if (cmd == NULL)
            break;
    }
    return ndyld;
}

uint32_t mask_opcode(uint32_t op) {
    return (op - (op & 0x0f));
}

uintptr_t read_uleb128(uintptr_t i, unsigned char *sym) {
    do {
        i++;
    } while(sym[i-1] & 0x80);
    i++;
    return i;
}

uintptr_t get_dyld_index(unsigned char *sym) {
    uintptr_t dyld_index = 0;

    for(uintptr_t i = 0; i < strlen((char *)sym); i++) {
        if(mask_opcode(sym[i]) == 0x40) {
            if(mask_opcode(sym[i-1]) == 0x90) {
                dyld_index = 0;
                break;
            }
            if(sym[i-2] == 0x20) {
                dyld_index = sym[i-1];
                break;
            }
            if(mask_opcode(sym[i-1]) == 0x10) {
                dyld_index = sym[i-1] - mask_opcode(sym[i-1]);
                break;
            }
        }
    }
    return dyld_index;
}

uintptr_t get_sym_name_start(unsigned char *sym) {
    uintptr_t i;
    for(i = 0; i < strlen((char*)sym); i++) {
        switch(sym[i]) {
            case 0x80: {
                i = read_uleb128(i, sym);
                break;
            }
            case 0xA0: {
                i = read_uleb128(i, sym);
                break;
            }
        }
        if(sym[i] == 0x40 || sym[i] == 0x41) {
            i++;
            break;
        }
    }
    return i;
}

void get_sym_name(unsigned char *sym) {
    uintptr_t start_index = get_sym_name_start(sym);
    if(start_index) {
        int q = 0;
        for(uintptr_t i = start_index; i < strlen((char*)sym); i++) {
            sym[q] = sym[i];
            q++;
        }
        sym[q] = '';
    }
}

void free_dyld_array(struct _dyld *dylds) {
    size_t arr_size = sizeof(dylds)/sizeof(dylds[0]);
    for(int i = 0; i < arr_size; i++) {
        free((void *)dylds[i].dyld_cmd);
        free(dylds[i].name);
    }
    free(dylds);
}

void free_dyld_cmd(void **dyld_cmd) {
    if(dyld_cmd != NULL) {
        free(*dyld_cmd);
        *dyld_cmd = NULL;
    }
}

/* Parse a Mach-O header */
void parse_macho (macho_input_t *input) {
    //Grab the entire header
    const struct mach_header_64 *header = macho_read(input, input->data, sizeof(*header));
    //Grab first load command
    const struct load_command *cmd = macho_offset(input, header, sizeof(*header), sizeof(struct load_command));

    //Init the array to store the dyld related load commands
    uintptr_t ndyld = get_dyld_count(input, header, cmd);
    if(!ndyld) {
        warnx("This binary does not seem to use any external frameworks/dylibs");
    }

    struct _dyld *dylds = malloc(sizeof(struct _dyld) * ndyld);
    uintptr_t ncmds = header->ncmds;

    const struct dyld_info_command *dyld_cmd = malloc(sizeof(struct dyld_info_command));
    uint32_t bind_off = 0;
    uint32_t bind_size = 0;
    const void *name_ptr = NULL;
    uint32_t placehold = 0;

    // Iterate over the load commands
    for (uintptr_t i = 0; i < ncmds; i++) {
        uint32_t cmdsize = cmd->cmdsize;
        cmd = macho_read(input, cmd, cmdsize);
        uint32_t cmd_type = cmd->cmd;
        switch (cmd_type) {
            case LC_DYLD_INFO_ONLY: {
                dyld_cmd = (const struct dyld_info_command *)cmd;
                bind_off = dyld_cmd->bind_off;
                bind_size = dyld_cmd->bind_size;
                break;
            }
            case LC_ID_DYLIB:
                case LC_LOAD_WEAK_DYLIB:
                case LC_REEXPORT_DYLIB:
                case LC_LOAD_DYLIB: {
                    dylds[placehold].dyld_cmd = (const struct dyld_info_command *)cmd;
                    size_t name_len = cmdsize - sizeof(struct dylib_command);
                    name_ptr = macho_offset(input, dylds[placehold].dyld_cmd, sizeof(struct dylib_command), name_len);
                    dylds[placehold].name = malloc(name_len);
                    memcpy(dylds[placehold].name, name_ptr, name_len);
                    placehold++;
                    break;
                }
        }
        // Load next command
        cmd = macho_offset(input, cmd, cmdsize, sizeof(struct load_command));
    }

    const void *pos_pointer = macho_offset(input, input->data, bind_off, 100);
    uintptr_t dyld_hold = 0;
    uint32_t offset_counter = 0;
    unsigned char *sym = malloc(1024);
    uintptr_t i = 1;

    while(offset_counter < (bind_size - 4)) {
        memcpy(sym, pos_pointer, 1024);
        pos_pointer+= (strlen((char*)sym) + 1);
        offset_counter += (strlen((char*)sym) + 1);
        uintptr_t dyld_index = get_dyld_index(sym);

        if(dyld_index) {
            dyld_hold = dyld_index;
        } else {
            dyld_index = dyld_hold;
        }

        get_sym_name(sym);
        if(strlen((char *)sym)) {
            printf("%lu : %s : %sn", i, sym, dylds[dyld_index].name);
        }

        i++;
    }

    free(sym);
    free_dyld_array(dylds);
}

int main (int argc, char *argv[]) {
    const char *path = argv[1];
    int fd = open(path, O_RDONLY);

    struct stat stbuf;
    stat(path, &stbuf);

    void *data = mmap(NULL, stbuf.st_size, PROT_READ, MAP_FILE|MAP_PRIVATE, fd, 0);

    macho_input_t input_file;
    input_file.data = data;
    input_file.length = stbuf.st_size;

    parse_macho(&input_file);

    munmap(data, stbuf.st_size);
    close(fd);
    printf("Donen");
    exit(0);
}