#define _XOPEN_SOURCE 600
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <stdbool.h>
#include <stdint.h>
#include <inttypes.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <sys/ptrace.h>
#include <sys/personality.h>
#include <sys/user.h>

#include "linenoise.h"
#include "mujs.h"
#include "hash.h"
#include "pmparser.h"
#include "da.h"
#include "libelfin_wrap.h"

#define LOG_ERR(fmt, ...) fprintf(stderr, "Error: " fmt, ##__VA_ARGS__)
#define LOG_INF(fmt, ...) fprintf(stdout, "Info:  " fmt, ##__VA_ARGS__)

#define unused(x) ((void)(x))

#define shift(argc, argv, msg, ...) (argc <= 0 ? (LOG_ERR(msg, ##__VA_ARGS__), exit(EXIT_FAILURE)) : (void)0, argc--, *argv++)

#define INIT_SCRIPT ".debugusrc.js"
#define MAX_JS_FUNCS 100

// Registers
// Taken from da goat https://source.winehq.org/source/dlls/dbghelp/cpu_x86_64.c

typedef enum {
    rax,
    rbx,
    rcx,
    rdx,
    rdi,
    rsi,
    rbp,
    rsp,
    r8,
    r9,
    r10,
    r11,
    r12,
    r13,
    r14,
    r15,
    rip,
    eflags,
    cs,
    fs,
    gs,
    ss,
    ds,
    es,
    tr,
    ldtr,
    mxcsr,
    ctrl,
    stat,
    MAX_REGISTERS,
} Register;

typedef struct {
    Register r;
    int dwarf_r;
    const char *name;
} RegisterDescriptor;

static RegisterDescriptor reg_descriptors[MAX_REGISTERS] = {
    #define make_rd(r, dr) { r, dr, #r }
    make_rd(rax, 0),
    make_rd(rdx, 1),
    make_rd(rcx, 2),
    make_rd(rbx, 3),
    make_rd(rsi, 4),
    make_rd(rdi, 5),
    make_rd(rbp, 6),
    make_rd(rsp, 7),
    make_rd(r8, 8),
    make_rd(r9, 9),
    make_rd(r10, 10),
    make_rd(r11, 11),
    make_rd(r12, 12),
    make_rd(r13, 13),
    make_rd(r14, 14),
    make_rd(r15, 15),
    make_rd(rip, 16),
    make_rd(eflags, 49),
    make_rd(es, 50),
    make_rd(cs, 51),
    make_rd(ss, 52),
    make_rd(ds, 53),
    make_rd(fs, 54),
    make_rd(gs, 55),
    make_rd(tr, 62),
    make_rd(ldtr, 63),
    make_rd(mxcsr, 64),
    make_rd(ctrl, 65),
    make_rd(stat, 66),
    #undef make_rd
};

// How breakpoints work?
// We can enable/disable breakpoints by putting/removing an int 3 instruction
// into/from the executed program. int 3 will trigger a SIGTRAP, which we can
// then handle. Pretty cool, eh?
//
// To handle a SIGTRAP, we can use waitpid(2) to listen for child process' events.

#define MAX_BRKS 0xff

typedef struct {
    pid_t pid;
    pid_t proc_pid;
    uintptr_t addr;
    bool enabled;
    uint8_t data;
} Brk;

void brk_enable(Brk *brk)
{
    long data = ptrace(PTRACE_PEEKDATA, brk->pid, brk->addr, NULL);
    brk->data = (uint8_t)(data&0xff);
    uint64_t int3 = 0xCC; // int 3 instruction, SIGTRAP
    uint64_t data_with_int3 = ((data& ~0xff) | int3);
    ptrace(PTRACE_POKEDATA, brk->pid, brk->addr, data_with_int3);

    brk->enabled = true;
}

void brk_disable(Brk *brk)
{
    long data = ptrace(PTRACE_PEEKDATA, brk->pid, brk->addr, NULL);
    long restored_data = ((data& ~0xff) | brk->data);
    ptrace(PTRACE_POKEDATA, brk->pid, brk->addr, restored_data);

    brk->enabled = false;
}

void symbols_deinit(Symbols *s)
{
    for (int i = 0; i < s->count; i++) {
        free((char *)s->items[i].name);
    }
    da_deinit(s);
}

typedef struct {
    const char *file;
    pid_t pid;
    js_State *js;
    HashTable brks;
    uintptr_t program_load_offset;
    HashTable js_descs;
    Symbols symbols;
    FILE *binfile;
    PLibelfinBinding plibelfin;
} Dbg;

siginfo_t dbg_get_siginfo(Dbg *dbg)
{
    siginfo_t i;
    ptrace(PTRACE_GETSIGINFO, dbg->pid, NULL, &i);
    return i;
}

void dbg_handle_sigsegv(Dbg *dbg, siginfo_t info)
{
    unused(dbg);
    uint64_t dbg_get_rip(Dbg *dbg);
    LOG_ERR("Caught a segfault %d. SKILL ISSUE BRO\n", info.si_code);
    uintptr_t rip = (uintptr_t)dbg_get_rip(dbg);
    uintptr_t nearest = rip - dbg->symbols.items[0].addr;
    for (int i = 0; i < dbg->symbols.count; i++) {
        if (rip - dbg->symbols.items[i].addr < nearest) {
            nearest = rip - dbg->symbols.items[i].addr;
        }
    }
    uintptr_t addr = rip - nearest;
    for (int i = 0; i < dbg->symbols.count; i++) {
        if (dbg->symbols.items[i].addr == addr) {
            AddrInfo *ai = libelfin_wrap_info_from_rip(dbg->plibelfin, (uint64_t)(addr - dbg->program_load_offset));
            if (ai != NULL) {
                LOG_INF("%s:%zu in function %s()\n", ai->file, (size_t)ai->line, dbg->symbols.items[i].name);
                libelfin_wrap_free_info(ai);
            }
            break;
        }
    }
}

void dbg_handle_sigtrap(Dbg *dbg, siginfo_t info)
{
    void dbg_set_rip(Dbg *dbg, uint64_t v);
    uint64_t dbg_get_rip(Dbg *dbg);

    AddrInfo *ai;

    switch (info.si_code) {
        case SI_KERNEL:
        case TRAP_BRKPT:
            dbg_set_rip(dbg, dbg_get_rip(dbg) - 1);
            ai = libelfin_wrap_info_from_rip(dbg->plibelfin, dbg_get_rip(dbg) - (uint64_t)dbg->program_load_offset);
            if (ai != NULL) {
                LOG_INF("Hit breakpoint at 0x%"PRIxPTR", %s:%zu\n", dbg_get_rip(dbg), ai->file, (size_t)ai->line);
                libelfin_wrap_free_info(ai);
            } else {
                LOG_INF("Hit breakpoint at 0x%"PRIxPTR"\n", dbg_get_rip(dbg));
            }
            return;
        case TRAP_TRACE:
            return;
    }
}

void dbg_wait(Dbg *dbg)
{
    int status, options = 0;
    waitpid(dbg->pid, &status, options);

    siginfo_t info = dbg_get_siginfo(dbg);
    switch (info.si_signo) {
        case SIGTRAP: dbg_handle_sigtrap(dbg, info); break;
        case SIGSEGV: dbg_handle_sigsegv(dbg, info); break;
        default:      LOG_INF("Signal %d\n!!", info.si_signo); break;
    }
}

// Memory

uint64_t dbg_mem_read(Dbg *dbg, uint64_t addr)
{
    return ptrace(PTRACE_PEEKDATA, dbg->pid, addr, NULL);
}

void dbg_mem_write(Dbg *dbg, uint64_t addr, uint64_t v)
{
    ptrace(PTRACE_POKEDATA, dbg->pid, addr, v);
}

uint64_t dbg_get_reg_value(Dbg *dbg, Register r)
{
    struct user_regs_struct rs;
    ptrace(PTRACE_GETREGS, dbg->pid, NULL, &rs);

    for (int i = 0; i < sizeof(reg_descriptors)/sizeof(reg_descriptors[0]); i++) {
        if (reg_descriptors[i].r == r) {
            return ((uint64_t*)&rs)[i];
        }
    }
}

void dbg_set_reg_value(Dbg *dbg, Register r, uint64_t v)
{
    struct user_regs_struct rs;
    ptrace(PTRACE_GETREGS, dbg->pid, NULL, &rs);

    for (int i = 0; i < sizeof(reg_descriptors)/sizeof(reg_descriptors[0]); i++) {
        if (reg_descriptors[i].r == r) {
            ((uint64_t*)&rs)[i] = v;
            ptrace(PTRACE_SETREGS, dbg->pid, NULL, &rs);
            return;
        }
    }
}

const char *get_reg_name(Register r)
{
    for (int i = 0; i < sizeof(reg_descriptors)/sizeof(reg_descriptors[0]); i++) {
        if (reg_descriptors[i].r == r) {
            return reg_descriptors[i].name;
        }
    }
}

Register get_reg_from_name(const char *name)
{
    for (int i = 0; i < MAX_REGISTERS; i++) {
        if (strcmp(name, reg_descriptors[i].name) == 0) {
            return reg_descriptors[i].r;
        }
    }
}

uint64_t dbg_get_rip(Dbg *dbg)
{
    return dbg_get_reg_value(dbg, rip);
}

void dbg_set_rip(Dbg *dbg, uint64_t v)
{
    dbg_set_reg_value(dbg, rip, v);
}

void dbg_step_brk(Dbg *dbg)
{
    uint64_t loc = dbg_get_rip(dbg);
    char key[20];
    snprintf(key, sizeof(key), "0x%"PRIxPTR, (uintptr_t)loc);
    Brk *brk = hashtable_get(&dbg->brks, key);
    if ((brk != NULL  && brk->enabled)) {
        brk_disable(brk);
        ptrace(PTRACE_SINGLESTEP, brk->pid, NULL, NULL);
        dbg_wait(dbg);
        brk_enable(brk);
    }
}

#define getdbg() ((Dbg*)js_currentfunctiondata(js))

void dbg_js_cont(js_State *js)
{
    Dbg *dbg = getdbg();
    dbg_step_brk(dbg);
    ptrace(PTRACE_CONT, dbg->pid, NULL, NULL);
    dbg_wait(dbg);

    js_pushundefined(js);
}

void dbg_js_baddr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    uintptr_t full_addr = dbg->program_load_offset + addr;
    char addr_str2[20];
    snprintf(addr_str2, sizeof(addr_str2), "0x%"PRIxPTR, full_addr);
    Brk brk = { .pid = dbg->pid, .addr = dbg->program_load_offset + addr };
    brk_enable(&brk);
    hashtable_set(&dbg->brks, addr_str2, &brk, sizeof(brk));
    js_pushundefined(js);
}

void dbg_js_bfn(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *fn_name = js_tostring(js, 1);
    for (int i = 0; i < dbg->symbols.count; i++) {
        Symbol *s = &dbg->symbols.items[i];
        if (strcmp(fn_name, s->name) == 0) {
            uintptr_t addr = s->addr;
            char addr_str2[20];
            snprintf(addr_str2, sizeof(addr_str2), "0x%"PRIxPTR, addr);
            Brk brk = { .pid = dbg->pid, .addr = addr };
            brk_enable(&brk);
            hashtable_set(&dbg->brks, addr_str2, &brk, sizeof(brk));
            break;
        }
    }
    js_pushundefined(js);
}

void dbg_js_rmbfn(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *fn_name = js_tostring(js, 1);

    for (int i = 0; i < dbg->symbols.count; i++) {
        Symbol *s = &dbg->symbols.items[i];
        if (strcmp(fn_name, s->name) == 0) {
            uintptr_t addr = s->addr;
            char addr_str2[20];
            snprintf(addr_str2, sizeof(addr_str2), "0x%"PRIxPTR, addr);
    
            Brk *brk = (Brk *)hashtable_get(&dbg->brks, addr_str2);
            if (brk == NULL) {
                LOG_ERR("No breakpoint at function: %s\n", fn_name);
                goto done;
            }
            brk_disable(brk);
            hashtable_delete(&dbg->brks, addr_str2);
            break;
        }
    }

done:
    js_pushundefined(js);
}

void dbg_js_rmbaddr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    Brk *brk = (Brk *)hashtable_get(&dbg->brks, addr_str);
    if (brk == NULL) {
        LOG_ERR("No breakpoint at address: %s\n", addr_str);
        goto done;
    }
    brk_disable(brk);
    hashtable_delete(&dbg->brks, addr_str);
done:
    js_pushundefined(js);
}

void dbg_js_splo(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    dbg->program_load_offset = addr;

    js_pushundefined(js);
}

void dbg_js_plo(js_State *js)
{
    Dbg *dbg = getdbg();
    char buf[20];
    snprintf(buf, sizeof(buf), "0x%"PRIxPTR, dbg->program_load_offset);
    js_pushstring(js, buf);
}

void dbg_js_lsbrk(js_State *js)
{
    Dbg *dbg = getdbg();
    int c = 0;
    for (int i = 0; i < dbg->brks.capacity; i++) {
        if (dbg->brks.buckets[i] != NULL) {
            Brk *brk = (Brk*)dbg->brks.buckets[i]->value;

            AddrInfo *ai = libelfin_wrap_info_from_rip(dbg->plibelfin, brk->addr - dbg->program_load_offset);
            if (ai != NULL) {
                LOG_INF("Breakpoint %d at 0x%"PRIxPTR" (%s) %s:%zu\n",
                        c+1, brk->addr, brk->enabled ? "Enabled" : "Disabled",
                        ai->file, (size_t)ai->line);
                libelfin_wrap_free_info(ai);
            } else {
                LOG_INF("Breakpoint %d at 0x%"PRIxPTR" (%s)\n",
                        c+1, brk->addr, brk->enabled ? "Enabled" : "Disabled");
            }
            c++;
        }
    }

    js_pushundefined(js);
}

void dbg_js_lsf(js_State *js)
{
    Dbg *dbg = getdbg();
    for (int i = 0; i < dbg->symbols.count; i++) {
        Symbol *s = &dbg->symbols.items[i];
        AddrInfo *ai = libelfin_wrap_info_from_rip(dbg->plibelfin, (uint64_t)(s->addr - dbg->program_load_offset));
        if (ai != NULL) {
            LOG_INF("Sym %s 0x%"PRIxPTR" %s:%zu\n", s->name, s->addr, ai->file, (size_t)ai->line);
            libelfin_wrap_free_info(ai);
        }
    }

    js_pushundefined(js);
}

void dbg_js_lif(js_State *js)
{
    const char *str = js_tostring(js, 1);
    LOG_INF("%s\n", str);
    js_pushundefined(js);
}

void dbg_js_ler(js_State *js)
{
    const char *str = js_tostring(js, 1);
    LOG_ERR("%s\n", str);
    js_pushundefined(js);
}

void dbg_js_ldscr(js_State *js)
{
    void dbg_load_script(Dbg *dbg, const char *script_path);
    Dbg *dbg = getdbg();
    const char *path = js_tostring(js, 1);
    dbg_load_script(dbg, path);
    js_pushundefined(js);
}

void dbg_js_gf(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushstring(js, dbg->file);
}

void dbg_js_pid(js_State *js)
{
    Dbg *dbg = getdbg();
    js_pushnumber(js, dbg->pid);
}

void dbg_js_gr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *name = js_tostring(js, 1);
    Register r = get_reg_from_name(name);
    uint64_t v = dbg_get_reg_value(dbg, r);
    char buf[20];
    snprintf(buf, sizeof(buf), "0x%"PRIx64, v);
    js_pushstring(js, buf);
}

void dbg_js_sr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *name = js_tostring(js, 1);
    Register r = get_reg_from_name(name);
    const char *value_str = js_tostring(js, 2);
    uint64_t value;
    sscanf(value_str, "0x%"SCNx64, &value);
    dbg_set_reg_value(dbg, r, value);
    js_pushundefined(js);
}

void dbg_js_mrd(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    uint64_t v = dbg_mem_read(dbg, addr);
    char buf[20];
    snprintf(buf, sizeof(buf), "0x%"PRIx64, v);
    js_pushstring(js, buf);
}

void dbg_js_mwr(js_State *js)
{
    Dbg *dbg = getdbg();
    const char *addr_str = js_tostring(js, 1);
    uintptr_t addr;
    sscanf(addr_str, "0x%"SCNxPTR, &addr);
    const char *value_str = js_tostring(js, 2);
    uint64_t value;
    sscanf(value_str, "0x%"SCNx64, &value);
    dbg_mem_write(dbg, addr, value);
    js_pushundefined(js);
}

void dbg_js_help(js_State *js)
{
    Dbg *dbg = getdbg();
    for (int i = 0; i < dbg->js_descs.capacity; i++) {
        if (dbg->js_descs.buckets[i] != NULL) {
            const char *key = dbg->js_descs.buckets[i]->key;
            char *val = dbg->js_descs.buckets[i]->value;
            LOG_INF("%s: %s\n", key, val);
        }
    }
    js_pushundefined(js);
}

void dbg_init_js(Dbg *dbg)
{
    dbg->js   = js_newstate(NULL, NULL, JS_STRICT);

    #define make_js_func(name, n, desc) \
    do { \
        js_newcfunctionx(dbg->js, &dbg_js_##name, #name, n, dbg, NULL); \
        js_setglobal(dbg->js, #name); \
        hashtable_set(&dbg->js_descs, #name, desc, strlen((desc))+1); \
    } while(0)

    make_js_func(cont, 0, "Continue execution, ARGS=None");
    make_js_func(baddr, 1, "Place breakpoint at address, ARGS=Address:hex string");
    make_js_func(rmbaddr, 1, "Remove breakpoint at address, ARGS=Address:hex string");
    make_js_func(lsbrk, 0, "List all breakpoints");
    make_js_func(splo, 1, "Set program load offset, ARGS=Offset:hex string");
    make_js_func(lif, 1, "Log information, ARGS=Text:string");
    make_js_func(ler, 1, "Log error, ARGS=Text:string");
    make_js_func(ldscr, 1, "Load script, ARGS=Path:string");
    make_js_func(gf, 0, "Get current executable file path, ARGS=None");
    make_js_func(pid, 0, "Get executable process ID, ARGS=None");
    make_js_func(plo, 0, "Get program load offset, ARGS=None");
    make_js_func(gr, 1, "Get value of a register, ARGS=Register name:string");
    make_js_func(sr, 2, "Set register value, ARGS=Register name:string,Register value:hex string");
    make_js_func(mrd, 1, "Read memory at address, ARGS=Address:hex string");
    make_js_func(mwr, 2, "Write memory at address, ARGS=Address:hex string,Value:hex string");
    make_js_func(help, 0, "Print help information, ARGS=None");
    make_js_func(lsf, 0, "List functions in executable, ARGS=None");
    make_js_func(bfn, 1, "Set breakpoint at function, ARGS=Function name:string");
    make_js_func(rmbfn, 1, "Remove breakpoint at function, ARGS=Function name:string");

    #undef make_js_func
}

void dbg_init_load_offset(Dbg *dbg)
{
    procmaps_iterator maps_iter = {0};
    procmaps_error_t parser_err = PROCMAPS_SUCCESS;

    parser_err = pmparser_parse(dbg->pid, &maps_iter);
    if (parser_err) {
        LOG_ERR("Failed to parse /proc/%d/maps (%d)\n", dbg->pid, (int)parser_err);
        return;
    }

    // We only need the first one
    procmaps_struct *mem_region = NULL;
    mem_region = pmparser_next(&maps_iter);
    dbg->program_load_offset = (uintptr_t)mem_region->addr_start;

    pmparser_free(&maps_iter);
}

void dbg_load_script(Dbg *dbg, const char *script_path)
{
    FILE *script = fopen(script_path, "r");
    if (script == NULL) {
        return;
    }

    fseek(script, 0L, SEEK_END);
    long size = ftell(script);
    rewind(script);
    char *script_buf = malloc(size+1);
    fread(script_buf, size, 1, script);
    script_buf[size] = '\0';

    js_dostring(dbg->js, script_buf);

    free(script_buf);
    fclose(script);
}

void dbg_init_bin(Dbg *dbg)
{
    FILE *bin = fopen(dbg->file, "rb");
    if (bin == NULL) {
        LOG_ERR("could not open file %s: %s\n", dbg->file, strerror(errno));
        return;
    }
    dbg->binfile = bin;
}

void dbg_libelfin_wrap_init(Dbg *dbg)
{
    dbg->plibelfin = libelfin_wrap_get_binding(fileno(dbg->binfile), dbg->program_load_offset);
}

void dbg_load_symbols(Dbg *dbg)
{
    libelfin_wrap_get_syms(dbg->plibelfin, &dbg->symbols);
}

void dbg_init(Dbg *dbg, const char *file, pid_t pid)
{
    memset(dbg, 0, sizeof(*dbg));
    dbg->file = file;
    dbg->pid  = pid;
    hashtable_init(&dbg->js_descs, MAX_JS_FUNCS);
    dbg_init_js(dbg);
    dbg_init_load_offset(dbg);
    dbg_init_bin(dbg);
    dbg_libelfin_wrap_init(dbg);
    dbg_load_symbols(dbg);
    hashtable_init(&dbg->brks, MAX_BRKS);
    dbg_load_script(dbg, INIT_SCRIPT);
}

void dbg_deinit(Dbg *dbg)
{
    js_freestate(dbg->js);
    hashtable_deinit(&dbg->brks);
    hashtable_deinit(&dbg->js_descs);
    symbols_deinit(&dbg->symbols);
    libelfin_wrap_free_binding(dbg->plibelfin);
    fclose(dbg->binfile);
}

void dbg_loop(Dbg *dbg)
{
    dbg_wait(dbg);
    char *line = NULL;
    while ((line = linenoise("debugus # ")) != NULL) {
        js_dostring(dbg->js, line);
        linenoiseHistoryAdd(line);
        linenoiseFree(line);
    }
}

void start_debugging(const char *input_file_path)
{
    pid_t pid = fork();
    if (pid == 0) {
        personality(ADDR_NO_RANDOMIZE);
        ptrace(PTRACE_TRACEME, 0, NULL, NULL);
        execl(input_file_path, input_file_path, NULL);
    } else if (pid >= 1) {
        Dbg dbg;
        dbg_init(&dbg, input_file_path, pid);
        dbg_loop(&dbg);
        dbg_deinit(&dbg);
    }
}

int main(int argc, char ** argv)
{
    const char *program_name = shift(argc, argv, "");
    unused(program_name);

    const char *input_file_path = shift(argc, argv, "No input file provided\n");
    start_debugging(input_file_path);

    return 0;
}