#include <elf.h>
#include <errno.h>
#include <linux/user.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ptrace.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include "libdasm/libdasm.h"

#define GI(dest,src) get_instruction(dest, src, MODE_32)

int save_entry_hook;
int hook_location;
int main_entry;
char entered_main;

Elf32_Ehdr elf;
Elf32_Phdr *phdrs;
Elf32_Sym *symtab;
char *strtab;

#define PAGEMASK 0xfffff000
#define PAGESIZE 0x1000

struct segment
{
    int start;
    int end;
    struct segment *next;
} segment;

struct func
{
    struct func *next;
    void *addr;
    char *name;
} func;

struct segment *seghead;
struct func *funchead;

void get_plt(int);
int ismapped(int);
char *lastfunc;


struct func *find_func_by_addr(int *addr)
{
    struct func *f;

    for(f = funchead; f; f = f->next)
    {
        if(f->addr == addr)
            return f;
    }
    return NULL;
}

int add_func(int *addr, char *name)
{
    struct func *f;

    if( find_func_by_addr(addr) ) return 0;

    f = calloc(1, sizeof(struct func));
    if(!f){ printf("Uhoh malloc: %m\n"); exit(0);}

    if( funchead )
        f->next = funchead;
    funchead = f;

    f->addr = addr;
    f->name = name;
    return 1;
}


int exec_program(char **args)
{
    int p;
    switch(p=fork())
    {
        case -1:
            perror("fork");
            return -1;
        case 0:
            ptrace(PTRACE_TRACEME, 0, 0, 0);
            execvp(args[0], args);
            perror("execv failed");
            kill(getppid(),9);
            exit(-1);
        default:
            wait(NULL); //...
            return p;
    }
}

int prevp;
int parse_position(int pid, int *next_eip)
{
    struct user_regs_struct regs;
    int ops[4];

    INSTRUCTION inst;    
    char string[256];   

    int x = 0;

    memset(&regs, 0, sizeof regs);
    ptrace(PTRACE_GETREGS, pid, 0, &regs);

    x = prevp == regs.eip;

    prevp = regs.eip;
    if( !ismapped(regs.eip) && !x)
         return regs.eip;

    if( regs.eip == main_entry )
        entered_main = 1;

    printf("eax=%.8x ebx=%.8x ecx=%.8x edx=%.8x "
           "esi=%.8x edi=%.8x ebp=%.8x esp=%.8x\n",
//           "eip=%.8x\n",
            regs.eax, regs.ebx, regs.ecx, regs.edx, 
            regs.esi, regs.edi, regs.ebp, regs.esp);
//            regs.eip);

//TODO, optimize me if necessary, corner cases on end of mapped pages?
    ops[0] = ptrace(PTRACE_PEEKTEXT, pid, regs.eip, 0);
    ops[1] = ptrace(PTRACE_PEEKTEXT, pid, regs.eip+4, 0);
    ops[2] = ptrace(PTRACE_PEEKTEXT, pid, regs.eip+8, 0);
    ops[3] = ptrace(PTRACE_PEEKTEXT, pid, regs.eip+12, 0);

    GI( &inst, &ops );
    get_instruction_string(&inst, FORMAT_ATT, regs.eip, string, sizeof(string));

//TODO: resove memory addresses, relevant vars
//  variable tracking ( to determine whats a string, whats a struct, etc )


    printf("%x: %s",regs.eip, string);

    if( inst.type == INSTRUCTION_TYPE_CALL || 
        inst.type == INSTRUCTION_TYPE_JMP ||
        inst.type == INSTRUCTION_TYPE_JMPC)
    {
        if( inst.immbytes )
        {
            int *dest =  (int *)((unsigned int)regs.eip + inst.op1.immediate + inst.length);
            struct func *f = NULL;
            f = find_func_by_addr(dest);
            if(f){
                printf(" \t###\t%s()",f->name);
                lastfunc = f->name;
            }
                        
        }
    }

    printf("\n");
////


    //TODO: assumption is the brother of failure
    if( inst.type == INSTRUCTION_TYPE_CALL )
        *next_eip = regs.eip + inst.length;
    else if (inst.type == INSTRUCTION_TYPE_RET)
        *next_eip = 0;


    return 0;

}


void set_hook(int pid, int location)
{
    hook_location = location; //save to global for run_hook

    //save dword @ location
    save_entry_hook = ptrace(PTRACE_PEEKTEXT, pid, location, 0);
    
    //hook it w/ an int3
    ptrace(PTRACE_POKETEXT, pid, location, (save_entry_hook & 0xffffff00)+0xcc);
}

void run_hook(int pid)
{
    struct user_regs_struct regs;
    int ret;
    ret = ptrace(PTRACE_CONT, pid, 0, 0);
    if(ret == -1) perror("ptrace");
    wait(NULL);
    printf("done waiting\n");

    //restore hook
    ret = ptrace(PTRACE_POKETEXT, pid, hook_location, save_entry_hook);
    if(ret == -1) {
      if(strcmp(lastfunc,"exit") == 0)
        exit(0);
       perror("ptrace on run_hook after PTRACE_CONT"); exit(1); 
    }

    //set eip to eip - 1
    ptrace(PTRACE_GETREGS, pid, 0, &regs);
    regs.eip -= 1;
    ptrace(PTRACE_SETREGS, pid, 0, &regs);
}


void walk_to_end(int pid)
{
    int ret,status=0;

    int next_eip; //used to set up hooks, dont want to single-step lib functions...
    int i = 0;
    do
    {
        ret = ptrace(PTRACE_SINGLESTEP, pid, 0, 0);

        wait(&status);

        if(ret == -1)
            perror("ptrace");
    
        ret = parse_position(pid, &next_eip);
        i += 1;
        if(ret != 0 && next_eip && entered_main ) //left to a segment that was not selected
        {

            set_hook(pid, next_eip);
            run_hook(pid);
        }
    
    
    } while (WIFSTOPPED(status));

    ptrace(PTRACE_CONT, pid, 0, 0);

}

int read_data(int pid, void *dest, int src, int len)
{
    int ret;
    int i = 0;
    unsigned int *p = (unsigned int *)dest;

    while(i*4 < len)
    {
        ret = ptrace(PTRACE_PEEKTEXT, pid, src+i*4, 0);
        if(errno) {      perror("ptrace read_data"); exit(1); }

        p[i++] = ret;
    }
    return i*4;
}

int read_string(int pid, int addr)
{
    char *string;
    int i, ret, cursize;

    cursize = 4096;
    string = malloc(cursize);
    if(string == NULL ) { perror("malloc read_string"); exit(1);}
    //length is unknown...

    string[0] = '\0';
    while(1)
    {
        ret = ptrace(PTRACE_PEEKTEXT, pid, addr, 0);
        addr += 4;

        if(errno) {      perror("ptrace read_string"); exit(1); }
        
        if(i+4 >= cursize){
            if(cursize + 4096 < cursize) {
                printf("LOLLERCOPTERZ\n"); exit(1);
            }
            string = realloc(string, cursize + 4096);
            if(string == NULL ) { perror("realloc read_string"); exit(1);}            
        }


        strncat(string, &ret, 4);

        //check for NULL string terminator
        if( (ret & 0xff) == 0  || (ret>>8&0xff) == 0
             || (ret>>16&0xff) == 0 || (ret>>24&0xff) == 0)
            break;

    }
    
    return string;
}


void read_headers(int pid)
{
    char buf[256];
    int ret, x, addr;
    FILE *f;

    //read maps
    ret = snprintf(buf, sizeof buf, "/proc/%d/maps",pid);
    if( ret >= sizeof buf || ret < 0)
    { printf("Path too long. Truncated to %s. Quitting\n",buf); exit(1); }
  
    f = fopen(buf, "r");
    if(!f){perror("maps"); exit(1);}
    fscanf(f, "%x-", &addr);
    fclose(f);

    
    read_data(pid, (void *)&elf, addr, sizeof(Elf32_Ehdr));

    phdrs = calloc(elf.e_phnum, sizeof(Elf32_Phdr));    
    read_data(pid, phdrs, addr + elf.e_phoff, sizeof(Elf32_Phdr)*elf.e_phnum);

    for(x = 0; x < elf.e_phnum; x++)
    {
        if(phdrs[x].p_type == PT_LOAD){
            
            struct segment *s = calloc(1, sizeof(struct segment));
            if(!s){ perror("Uhoh malloc segment\n"); exit(1);}

            if(seghead)
                s->next = seghead;
            seghead = s;
            seghead->start = phdrs[x].p_vaddr & PAGEMASK;
            seghead->end = ((phdrs[x].p_vaddr + phdrs[x].p_memsz) & PAGEMASK) + PAGESIZE;
                
        }
    }

}

void hook_entry(int pid)
{

/*
    //fetch ehdr->e_entry;
    ret = ptrace(PTRACE_PEEKTEXT, pid, &(e->e_entry), 0);
    if(errno) {      perror("ptrace"); exit(1); }
*/
    read_headers(pid);


    set_hook(pid, elf.e_entry);

    main_entry = ptrace(PTRACE_PEEKTEXT, pid, elf.e_entry+24, 0);
    if(errno) {      perror("ptrace find main"); exit(1); }

}

int main(int argc, char **argv)
{
    int pid;

    if(!argv[1])
    {   printf("Usage: %s target args\n",argv[0]);   exit(0); }

    pid = exec_program(argv+1);

    if(pid == -1)
    { printf("FAILED TO EXECUTE TARGET %s\n",argv[1]);  exit(0); }

    hook_entry(pid);

    get_plt(pid);

    run_hook(pid);
    walk_to_end(pid);

    return 0;
}







int ismapped(int a)
{
    struct segment *s;
    s = seghead;
    while(s)
    {
        if((unsigned int)a <= s->end && (unsigned int)a >= s->start)
            return 1;
        s = s->next;
    }
    return 0;
}


char *lookup_rel(int pid, Elf32_Rel *r)
{
    char type;
    uint pos;
    Elf32_Sym *s;
    Elf32_Rel relstruct;
    Elf32_Sym symstruct;

    read_data(pid, &relstruct, r, sizeof(Elf32_Rel));


    type = relstruct.r_info&0xff;
    pos = relstruct.r_info >> 8;

    s = &symtab[pos];

    read_data(pid, &symstruct, s, sizeof(Elf32_Sym));

    if(symstruct.st_name)
        return read_string(pid, symstruct.st_name+strtab);
    else
        return NULL;

}
void get_plt(int pid)
{
    int i, pltrelsz;
    int addr;
    Elf32_Dyn d;
    Elf32_Rel *r, *ri, *jr;
    Elf32_Rel relstruct;

    for(i = 0; i < elf.e_phnum; i++)
    {
        if( phdrs[i].p_type == PT_DYNAMIC )
            break;
    }
    if(i == elf.e_phnum)
    {
        printf("MISSING DYNAMIC PHDR\n");
        return;
    }

    addr = phdrs[i].p_vaddr;

    read_data(pid, &d, addr, sizeof(Elf32_Dyn));

    for( ; d.d_tag != DT_NULL; )
    {
        if(d.d_tag == DT_REL)
            r = d.d_un.d_val;
        else if(d.d_tag == DT_PLTRELSZ)
            pltrelsz = d.d_un.d_val;
        else if(d.d_tag == DT_JMPREL)
            jr = d.d_un.d_val;
        else if(d.d_tag == DT_SYMTAB)
            symtab = d.d_un.d_val;
        else if(d.d_tag == DT_STRTAB)
            strtab = d.d_un.d_ptr;

        addr += sizeof(Elf32_Dyn);
        read_data(pid, &d, addr, sizeof(Elf32_Dyn));
    }

    for( ri = jr ; (unsigned int)ri < (unsigned int)jr + pltrelsz; ri++ )
    {
        char *s;
        s = lookup_rel(pid, ri);

        if(s)
            s = strdup(s);
        else
            s = "unknown";
        int func;

//        func = (*(int *)( relstruct->r_offset))  -6;

        read_data(pid, &relstruct, ri, sizeof(Elf32_Rel));
        read_data(pid, &func, relstruct.r_offset, sizeof(int *));
        func -= 6;

//        printf("GOT@%x plt@%x %s\n",ri->r_offset,func,s);
        add_func(func, s);
    }

}