/*Some sort of disassembling frontend...
Author: Andreas Guðmundsson
using http://udis86.sourceforge.net/
gcc  -ludis86 disassembler.c -o disas  --std=c99
*/
#include <mach-o/loader.h>
#include <mach/i386/thread_status.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <udis86.h> 

#define debug printf

struct {
    uint8_t mode_bits;
    void (*syntax)(ud_t*);
    FILE *filep;
    enum {RAW, MACHO} filet; /*XXX PE32, PE32+, ELF*/
    struct {
	int maddr;
	int foffset;
    } entryp;
} config;
static const char *filet_str[] = {"raw", "mach-o"};
/*add option to specify from which address to disassemble, allow both file offset and inmemory address*/
void config_init(void) {
    config.mode_bits = 32;
    config.syntax = UD_SYN_INTEL;
    config.filep = stdin;
    config.filet = RAW;
    config.entryp.maddr = 0;
    config.entryp.foffset = 0;
}
void usage(void) {
    printf("Options:\n"
	   "-h       : prints this usage information\n"
	   "-m       : 16,32 or 64 bit mode (%d default)\n"
	   "-s       : intel or at(at&t) syntax (%s default)\n" 
	   "-t       : raw or macho(mach-o) (%s default)\n"
	   "-f       : input file (stdin default)\n",
	   config.mode_bits, 
	   (config.syntax == UD_SYN_INTEL ? "intel" : "at&t"),
	   filet_str[config.filet]
	);
}
void parse_argv(int argc, char **argv) {
#define invalid_arg(opt,arg,fmt,def) fprintf(stderr, "%s is an invalid argument to %s, defaulting to " fmt "\n", arg, opt, def)
    int c;
    while ((c = getopt(argc, argv, "hm:s:t:f:")) != -1) {
	switch (c) {
	case 'h':
	    usage();
	    exit(1);
	case 'm':
	    if (0 == strcmp("16", optarg)) {
		config.mode_bits = 16;
	    } else if (0 == strcmp("32", optarg)) {
		config.mode_bits = 32;
	    } else if (0 == strcmp("64", optarg)) {
		config.mode_bits = 64;
	    } else {
		invalid_arg("-m", optarg, "%d", config.mode_bits);
	    }
	    break;
	case 's':
	    if (0 == strcasecmp("intel", optarg)) {
		config.syntax = UD_SYN_INTEL;
	    } else if (0 == strcasecmp("at&t", optarg) || 0 == strcasecmp("att", optarg)) {
		config.syntax = UD_SYN_ATT;
	    } else {
		if (config.syntax == UD_SYN_INTEL) {
		    invalid_arg("-s", optarg, "%s", "intel");
		} else {
		    invalid_arg("-s", optarg, "%s", "at&t");
		}
	    }
	    break;
	case 't':
	    if (0 == strcasecmp("raw", optarg)) {
		config.filet = RAW;
	    } else if (0 == strcasecmp("macho", optarg) || 0 == strcasecmp("mach-o", optarg)) {
		config.filet = MACHO;
	    } else {
		switch(config.filet) {
		case RAW: invalid_arg("-t", optarg, "%s", "raw"); break;
		case MACHO: invalid_arg("-t", optarg, "%s", "mach-o"); break;
		}
	    }
	    break;
	case 'f':
	    config.filep = fopen(optarg, "r");
	    if (!config.filep) {
		perror("Aborting");
		exit(1);
	    }
	    break;
	case '?':
	    if (optopt == 'c')
		fprintf (stderr, "Option -%c requires an argument.\n", optopt);
	    else if (isprint (optopt))
		fprintf (stderr, "Unknown option `-%c'.\n", optopt);
	    else
		fprintf (stderr, "Unknown option character `\\x%x'.\n", optopt);
	default:
	    usage();
	    exit(1);
	}
    }
#undef invalid_arg
}
void macho_parse(void) {
	/*Getting tired... 
	  want to parse mach-o info, 
	  deal with universal binaries,
	  so I'm lazy right now and just going for the entrypoint of i386*/
    struct x86_thread_state x86ts;
    struct mach_header mh;
    struct load_command lc;
    fseek(config.filep, 0, SEEK_SET);
    fread(&mh, sizeof(mh), 1, config.filep);
    for (uint32_t i = 0; i < mh.ncmds; i++) {
	fread(&lc, sizeof(lc), 1, config.filep);
	switch(lc.cmd) {
	case LC_THREAD:
	case LC_UNIXTHREAD:
	    fread(&x86ts, sizeof(x86ts), 1, config.filep);
	    config.entryp.maddr = x86ts.uts.ts32.eip;
	    printf("Found entrypoint inmemory address 0x%x\n", config.entryp.maddr);
	    /*I should gather segment info to figure out the file offset*/
	    /*lets be lazy and loop through the load commands looking for the entryp,
	      this won't do much if the entryp is out of segment*/
	    long pos = ftell(config.filep);
	    fseek(config.filep, sizeof(mh), SEEK_SET);
	    printf("NCMDS %d\n", mh.ncmds);
	    for (uint32_t j = 0; j < mh.ncmds; j++) {
		fread(&lc, sizeof(lc), 1, config.filep);
		printf("CMD %d\n", lc.cmd);
		if (lc.cmd == LC_SEGMENT) {
		    struct segment_command sc;
		    fread((char *)&sc+sizeof(lc), sizeof(sc)-sizeof(lc), 1, config.filep);
		    printf("Looking in %s\n", sc.segname);
		    if (sc.vmaddr <= config.entryp.maddr && config.entryp.maddr < sc.vmaddr+sc.vmsize) {
			config.entryp.foffset = (config.entryp.maddr-sc.vmaddr)+sc.fileoff;
			printf("Found entrypoint file offset 0x%x\n", config.entryp.foffset);
			break;
		    }
		} else {
		    fseek(config.filep, lc.cmdsize-sizeof(lc), SEEK_CUR);
		}
	    }
	    /*restore filepointer*/
	    fseek(config.filep, pos, SEEK_SET);
	    if (config.entryp.foffset == 0) 
		    printf("Didn't find entrypoint file offset!\n");
	}
	fseek(config.filep, lc.cmdsize-sizeof(lc), SEEK_CUR);
    }
}
int main(int argc, char **argv) 
{ 
    ud_t ud_obj; 
    config_init();
    parse_argv(argc, argv);
/*  buffer the stream if config.filep == stdin */
    switch (config.filet) {
    case RAW:
	/*just disassemble online given the config,
	  I guess there won't be any work done here*/
	break;
    case MACHO: 
	macho_parse();
    }
    /*yes sir I need sleep*/
    /*for now lets just start disassembling from the entrypoint*/
    fseek(config.filep, config.entryp.foffset, SEEK_SET);
    ud_init(&ud_obj); 
    ud_set_input_file(&ud_obj, config.filep); 
    ud_set_mode(&ud_obj, config.mode_bits); 
    ud_set_syntax(&ud_obj, config.syntax); 
    ud_set_pc(&ud_obj, config.entryp.maddr);
    while (ud_disassemble(&ud_obj)) { 
	printf("\t%s\n", ud_insn_asm(&ud_obj)); 
    } 
    return 0; 
}