#include "dr_api.h"
#include "dr_ir_opnd.h"
#include "dr_ir_instr.h"
#include "dr_ir_opcodes_x86.h"
#include "dr_ir_macros_x86.h"

#include <ctype.h>
#include <string.h>
#include <windows.h>
#include <commdlg.h>

char *CUSTOM_FILE = "C:\\Users\\vboxuser\\Downloads\\spoofed_number.txt";

static app_pc pc_GetOpenFileNameA = NULL;

static void intercept_GetOpenFileNameA() {
    // dr_printf("intercept_GetOpenFileNameA(...)\n");

    void *drcontext = dr_get_current_drcontext();

    dr_mcontext_t mcontext = { sizeof(mcontext), DR_MC_ALL };
    dr_get_mcontext(drcontext, &mcontext);

    OPENFILENAME *ofn = (OPENFILENAME *)mcontext.rcx;
    // dr_printf("OPENFILENAME at %p\n", ofn);

    for (int i = 0; i < strlen(CUSTOM_FILE); i++) ofn->lpstrFile[i] = CUSTOM_FILE[i];
    ofn->lpstrFile[strlen(CUSTOM_FILE)] = 0;

    mcontext.rax = 1;

    dr_set_mcontext(drcontext, &mcontext);
}

static void event_exit(void) {
    // dr_printf("Exiting custom DynamoRIO client...\n");
}

static void event_module_load(void *drcontext, const module_data_t *mod, bool loaded) {
    char lowerpath[MAX_PATH];
    strcpy(lowerpath, mod->full_path);
    for (int i = 0; i < strlen(lowerpath); i++) {
        lowerpath[i] = tolower(lowerpath[i]);
    }

    if (strstr(lowerpath, "comdlg32.dll")) {
        pc_GetOpenFileNameA = (app_pc)dr_get_proc_address(mod->handle, "GetOpenFileNameA");
        // if (pc_GetOpenFileNameA) dr_printf("Found GetOpenFileNameA at %p\n", pc_GetOpenFileNameA);
    }
}


static dr_emit_flags_t event_basic_block(void *drcontext, void *tag, instrlist_t *bb, bool for_trace, bool translating) {
    instr_t *to_remove = NULL, *instr = instrlist_first(bb);
    while (instr) {
        instr_t *next = instr_get_next(instr);

        app_pc target = NULL;

        if (instr_is_call_direct(instr)) {
            app_pc target = instr_get_branch_target_pc(instr);
        } else if (instr_is_call_indirect(instr)) {
            opnd_t opnd = instr_get_src(instr, 0);
            if (opnd_is_memory_reference(opnd)) {
                app_pc addr = opnd_compute_address(opnd, drcontext);
                if (addr) dr_safe_read(addr, sizeof(app_pc), &target, NULL);
            }
        }

        if (target == pc_GetOpenFileNameA) {
            // dr_printf("Call to GetOpenFileNameA detected at %p\n", instr_get_app_pc(instr));

            dr_insert_clean_call(drcontext, bb, instr, (void *)intercept_GetOpenFileNameA, false, 0);

            instrlist_remove(bb, instr);
            instr_destroy(drcontext, instr);
            instr = next;
        }

        instr = next;
    }
    
    return DR_EMIT_DEFAULT;
}

DR_EXPORT void dr_client_main(client_id_t id, int argc, const char *argv[]) {
    // dr_printf("Loading custom DynamoRIO client...\n");
    dr_set_client_name("Custom DynamoRIO client", "https://gitea.cloud.lehnert.dev/ludwig/windows-binary-fuzzing");
    dr_register_exit_event(event_exit);
    dr_register_module_load_event(event_module_load);
    dr_register_bb_event(event_basic_block);
    // dr_printf("Custom DynamoRIO client loaded.\n");
}