/*
 * hikaru.c
 *
 * Used for parsing a binary file and determining
 * possible instruction sequences.
 *
 * Requires libdisasm 0.21-pre2, and libelf 0.8.4
 *
 * Email all questions/comments to comsatcat at gmail.com
 *
 */

#include <stdio.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <sysexits.h>
#include <sys/mman.h>
#include <libelf.h>
#include <gelf.h>
#include <libdis.h>

#define TRUE    1
#define FALSE   2
#define M_ANALYZE       1
#define M_PROCESS       2
#define M_LIST          3

#define MAX_SIZE 32

void usage(char *);
int process(char *, int bytes, Elf *);
uint32_t getentrypoint(Elf *);
void listsections(Elf *);
void getsectioninfo(char *, Elf *);
void x86dis_manual_print( x86_insn_t *, void *);
void analyze(char *, Elf *);
int analyze_process(char *, Elf *);

struct {
        uint32_t offset;
        uint32_t size;
} si;

int fd;
int output = FALSE;
char outbuf[8192];

int main(int argc, char **argv) {
        int bytes = 0;
        int ch, n, list = 0, count_ret = 0, count_total = 0, mode = 0;
        char *file = NULL, *section, *image, buf[2];
        enum x86_options options = opt_none;
        Elf *elf;

        while((ch = getopt(argc, argv, "m:f:b:s:l")) != -1) {
                switch(ch) {
                        case 'm':
                                if(!strcasecmp(optarg, "analyze"))
                                        mode = M_ANALYZE;
                                else if(!strcasecmp(optarg, "process"))
                                        mode = M_PROCESS;
                        break;
                        case 'f':
                                if(!(file = malloc(strlen(optarg)+1))) {
                                        perror("malloc");
                                        exit(1);
                                }
                                strcpy(file, optarg);
                        break;
                        case 'b':
                                bytes = atoi(optarg);
                        break;
                        case 's':
                                section = malloc(strlen(optarg) + 1);
                                strcpy(section, optarg);
                        break;
                        case 'l':
                                list = TRUE;
                        break;
                }
        }

        if((!mode) || (!file) || ((mode == M_PROCESS) && (!bytes) && (!section)))
                usage(argv[0]);

        if((fd = open(file, O_RDONLY)) == -1) {
                perror("open");
                exit(1);
        }

        if(elf_version(EV_CURRENT) == EV_NONE) {
                errx(EX_SOFTWARE, "ELF library initialization failed: %s", elf_errmsg(-1));
        }

        if(!(elf = elf_begin(fd, ELF_C_READ, NULL))) {
                errx(EX_NOINPUT, "open \"%s\" failed", file);
        }

        if(elf_kind(elf) != ELF_K_ELF) {
                errx(EX_DATAERR, "%s is not an ELF object.", file);
        }

        if(mode == M_ANALYZE) {
                if(list) {
                        printf("Entry Point: 0x%x\n\n", getentrypoint(elf));
                        listsections(elf);
                } else {
                        analyze(file, elf);
                }
        } else if(mode == M_PROCESS) {
                x86_init(options, NULL);
                if(process(section, bytes, elf))
                        errx(EX_SOFTWARE, "failed to process file.");
        }

        close(fd);
        elf_end(elf);
}

void analyze(char *file, Elf *elf) {
        char *name;
        size_t shstrndx;
        Elf_Scn *scn = NULL;
        GElf_Shdr shdr;

        if(elf_getshstrndx(elf, &shstrndx)) {
                errx(EX_SOFTWARE, "getshstrndx failed: %s", elf_errmsg(-1));
        }

        while ((scn = elf_nextscn(elf, scn)) != NULL) {
                if (gelf_getshdr(scn, &shdr) != &shdr)
                        errx(EX_SOFTWARE, "getshdr failed: %s.", elf_errmsg(-1));
                if (!(name = elf_strptr(elf, shstrndx, shdr.sh_name)))
                        errx(EX_SOFTWARE, "elf_strptr failed: %s.", elf_errmsg(-1));

                if(shdr.sh_flags & PF_X) {
                        printf("Analyzing [%04jd] (%s) @ offset 0x%08x ", (uintmax_t) elf_ndxscn(scn), name, shdr.sh_offset);
                        printf("(%d bytes)\n", shdr.sh_size);
                        analyze_process(name, elf);
                }
        }
}

uint32_t getentrypoint(Elf *elf) {
        GElf_Phdr phdr;
        gelf_getphdr(elf, 0, &phdr);
        return(phdr.p_vaddr - phdr.p_offset);
}

void listsections(Elf *elf) {
        char *name, *tmp;
        size_t shstrndx;
        Elf_Scn *scn = NULL;
        GElf_Shdr shdr;

        if(elf_getshstrndx(elf, &shstrndx)) {
                errx(EX_SOFTWARE, "getshstrndx failed: %s", elf_errmsg(-1));
        }

        printf("Index\tOffset\t\tSize (bytes)\tSection\n");
        printf("-----\t-----------\t------------\t-------\n");
        while ((scn = elf_nextscn(elf, scn)) != NULL) {
                if (gelf_getshdr(scn, &shdr) != &shdr)
                        errx(EX_SOFTWARE, "getshdr failed: %s.", elf_errmsg(-1));
                if (!(name = elf_strptr(elf, shstrndx, shdr.sh_name)))
                        errx(EX_SOFTWARE, "elf_strptr failed: %s.", elf_errmsg(-1));
                printf("%04jd\t0x%08x\t", (uintmax_t) elf_ndxscn(scn), shdr.sh_offset);
                printf("%d\t\t", shdr.sh_size);
                printf("%s\n", name);
        }
        printf("\n");
}

void getsectioninfo(char *section, Elf *elf) {
        char *name;
        size_t shstrndx;
        Elf_Scn *scn = NULL;
        GElf_Shdr shdr;

        if(elf_getshstrndx(elf, &shstrndx)) {
                errx(EX_SOFTWARE, "getshstrndx failed: %s", elf_errmsg(-1));
        }
        while ((scn = elf_nextscn(elf, scn)) != NULL) {
                if (gelf_getshdr(scn, &shdr) != &shdr)
                        errx(EX_SOFTWARE, "getshdr failed: %s.", elf_errmsg(-1));
                if (!(name = elf_strptr(elf, shstrndx, shdr.sh_name)))
                        errx(EX_SOFTWARE, "elf_strptr failed: %s.", elf_errmsg(-1));

                if(!strcmp(name, section)) {
                        si.offset = shdr.sh_offset;
                        si.size = shdr.sh_size;
                }
        }
}

int process(char *section, int bytes, Elf *elf) {
        uint32_t entry;
        unsigned char *image, *buf;
        struct stat s;
        int i, t;

        if(fstat(fd, &s) == -1) {
                perror("fstat");
                return(1);
        }

        entry = getentrypoint(elf);
        getsectioninfo(section, elf);

        if(!si.offset && !si.size) {
                fprintf(stderr, "Section '%s' does not exist.\n", section);
                return(1);
        }

        if((int)(image = (unsigned char *)mmap(NULL, s.st_size, PROT_READ, MAP_SHARED, fd, 0)) == -1) { // fd is global
                fprintf(stderr, "Unable to map anonymous memory: %s\n", strerror(errno));
                return(1);
        }

        for(i=0; i<si.size; i++) {
                buf = image;
                x86_disasm_range(buf, entry, si.offset+i, bytes, x86dis_manual_print, NULL);
                memset(outbuf, 0, sizeof(outbuf));
        }

        return(0);
}

int analyze_process(char *section, Elf *elf) {
        uint32_t entry;
        unsigned char *image, *buf;
        struct stat s;
        int i, t;

        if(fstat(fd, &s) == -1) {
                perror("fstat");
                return(1);
        }

        entry = getentrypoint(elf);
        getsectioninfo(section, elf);

        if(!si.offset && !si.size) {
                fprintf(stderr, "Section '%s' does not exist.\n", section);
                return(1);
        }

        if((int)(image = (unsigned char *)mmap(NULL, s.st_size, PROT_READ, MAP_SHARED, fd, 0)) == -1) { // fd is global
                fprintf(stderr, "Unable to map anonymous memory: %s\n", strerror(errno));
                return(1);
        }

        for(i=0; i<si.size; i++) {
                for(t=MAX_SIZE; t>0; t--) {
                        buf = image;
                        if(t > si.size)
                                t = si.size;

                        x86_disasm_range(buf, entry, si.offset+i, t, x86dis_manual_print, NULL);
                        memset(outbuf, 0, sizeof(outbuf));
                }
        }

        return(0);
}

// this function was initially ripped from x86dis.c
// however its going to be heavily modified and work
// in a very _gross_ way.
void x86dis_manual_print( x86_insn_t *insn, void *arg ) {
        char line[4096];
        int i;

        if ( x86_format_insn(insn, line, 4096, att_syntax) <= 0 ) {
                return;
        }

        sprintf(outbuf, "%s%08X", outbuf, insn->addr);
        for ( i = 0; i < 10; i++ ) {
                if ( i < insn->size ) {
                        sprintf(outbuf, "%s %02X", outbuf, insn->bytes[i]);
                } else {
                        strcat(outbuf, "   ");
                }
        }

        if(!strcmp(line, "ret\t")) {
                sprintf(outbuf, "%s\t%s\n\n", outbuf, line);
                printf("%s", outbuf);
                memset(outbuf, 0, sizeof(outbuf));
        } else {
                sprintf(outbuf, "%s\t%s\n", outbuf, line);
        }
}

void usage(char *p) {
        printf("usage: %s <-f file> <[-m analyze [-l]] [-m process <-s section> <-b bytes>]>\n\n", p);
        printf("Hikaru operates in two modes, either enalyze mode or process mode.\n");
        printf("Both modes take a -f argument which defines a file to process.  \n\n");
        printf("Analyze mode has a single optional argument, -l, which will list \n");
        printf("the elf header sections available for processing.  If Analyze mode \n");
        printf("is run without the -l flag, it will analyze all executable sections \n");
        printf("of the target binary displaying all instructions up to MAX_SIZE.\n\n");
        printf("Process mode is similar to analyze, however it takes a specific \n");
        printf("section to process and the size of the instruction to analyze in bytes.\n\n");
        printf("As of this release, invalid operands are not displayed, however if a \n");
        printf("instruction sequence contains a ret and still has invalid operands \n");
        printf("before the ret instruction, all valid instructions up to the ret will \n");
        printf("be displayed.\n\n");
        exit(0);
}