#include "stack_walk.h"

#include "module.h"
#include "utils.h"
#include "win_funcs.h"

// =======================
// ====== windows.h ======
// =======================

#define CONTEXT_AMD64   0x00100000L

#define CONTEXT_CONTROL         (CONTEXT_AMD64 | 0x00000001L)
#define CONTEXT_INTEGER         (CONTEXT_AMD64 | 0x00000002L)
#define CONTEXT_SEGMENTS        (CONTEXT_AMD64 | 0x00000004L)
#define CONTEXT_FLOATING_POINT  (CONTEXT_AMD64 | 0x00000008L)
#define CONTEXT_DEBUG_REGISTERS (CONTEXT_AMD64 | 0x00000010L)

#define CONTEXT_FULL            (CONTEXT_CONTROL | CONTEXT_INTEGER | \
                                 CONTEXT_FLOATING_POINT)

#define CONTEXT_ALL             (CONTEXT_CONTROL | CONTEXT_INTEGER | \
                                 CONTEXT_SEGMENTS | CONTEXT_FLOATING_POINT | \
                                 CONTEXT_DEBUG_REGISTERS)

#define CONTEXT_XSTATE          (CONTEXT_AMD64 | 0x00000040L)
#define CONTEXT_KERNEL_CET      (CONTEXT_AMD64 | 0x00000080L)

#define PAGE_NOACCESS           0x01    
#define PAGE_READONLY           0x02    
#define PAGE_READWRITE          0x04    
#define PAGE_WRITECOPY          0x08    
#define PAGE_EXECUTE            0x10    
#define PAGE_EXECUTE_READ       0x20    
#define PAGE_EXECUTE_READWRITE  0x40    
#define PAGE_EXECUTE_WRITECOPY  0x80    
#define PAGE_GUARD             0x100    
#define PAGE_NOCACHE           0x200    
#define PAGE_WRITECOMBINE      0x400    
#define PAGE_GRAPHICS_NOACCESS           0x0800    
#define PAGE_GRAPHICS_READONLY           0x1000    
#define PAGE_GRAPHICS_READWRITE          0x2000    
#define PAGE_GRAPHICS_EXECUTE            0x4000    
#define PAGE_GRAPHICS_EXECUTE_READ       0x8000    
#define PAGE_GRAPHICS_EXECUTE_READWRITE 0x10000    
#define PAGE_GRAPHICS_COHERENT          0x20000    
#define PAGE_GRAPHICS_NOCACHE           0x40000    
#define PAGE_ENCLAVE_THREAD_CONTROL 0x80000000  
#define PAGE_REVERT_TO_FILE_MAP     0x80000000  
#define PAGE_TARGETS_NO_UPDATE      0x40000000  
#define PAGE_TARGETS_INVALID        0x40000000  
#define PAGE_ENCLAVE_UNVALIDATED    0x20000000  
#define PAGE_ENCLAVE_MASK           0x10000000  
#define PAGE_ENCLAVE_DECOMMIT       (PAGE_ENCLAVE_MASK | 0) 
#define PAGE_ENCLAVE_SS_FIRST       (PAGE_ENCLAVE_MASK | 1) 
#define PAGE_ENCLAVE_SS_REST        (PAGE_ENCLAVE_MASK | 2) 
#define MEM_COMMIT                      0x00001000  
#define MEM_RESERVE                     0x00002000  
#define MEM_REPLACE_PLACEHOLDER         0x00004000  
#define MEM_RESERVE_PLACEHOLDER         0x00040000  
#define MEM_RESET                       0x00080000  
#define MEM_TOP_DOWN                    0x00100000  
#define MEM_WRITE_WATCH                 0x00200000  
#define MEM_PHYSICAL                    0x00400000  
#define MEM_ROTATE                      0x00800000  
#define MEM_DIFFERENT_IMAGE_BASE_OK     0x00800000  
#define MEM_RESET_UNDO                  0x01000000  
#define MEM_LARGE_PAGES                 0x20000000  
#define MEM_4MB_PAGES                   0x80000000  
#define MEM_64K_PAGES                   (MEM_LARGE_PAGES | MEM_PHYSICAL)  
#define MEM_UNMAP_WITH_TRANSIENT_BOOST  0x00000001  
#define MEM_COALESCE_PLACEHOLDERS       0x00000001  
#define MEM_PRESERVE_PLACEHOLDER        0x00000002  
#define MEM_DECOMMIT                    0x00004000  
#define MEM_RELEASE                     0x00008000  
#define MEM_FREE                        0x00010000  
#define MEM_EXTENDED_PARAMETER_GRAPHICS                 0x00000001  
#define MEM_EXTENDED_PARAMETER_NONPAGED                 0x00000002  
#define MEM_EXTENDED_PARAMETER_ZERO_PAGES_OPTIONAL      0x00000004  
#define MEM_EXTENDED_PARAMETER_NONPAGED_LARGE           0x00000008  
#define MEM_EXTENDED_PARAMETER_NONPAGED_HUGE            0x00000010  
#define MEM_EXTENDED_PARAMETER_SOFT_FAULT_PAGES         0x00000020  
#define MEM_EXTENDED_PARAMETER_EC_CODE                  0x00000040  
#define MEM_EXTENDED_PARAMETER_NUMA_NODE_MANDATORY      MINLONG64	  

// ========================
// ====== !windows.h ======
// ========================

// enum UNWIND_INFO_FLAGS {
//     UNW_FLAG_EHANDLER,
//     UNW_FLAG_UHANDLER,
//     UNW_FLAG_CHAININFO,
// };

enum UNWIND_OP_TYPE {
    UWOP_PUSH_NONVOL,
    UWOP_ALLOC_LARGE,
    UWOP_ALLOC_SMALL,
    UWOP_SET_FPREG,
    UWOP_SAVE_NONVOL,
    UWOP_SAVE_NONVOL_FAR,
    UWOP_SAVE_XMM128 = 8U,
    UWOP_SAVE_XMM128_FAR,
    UWOP_PUSH_MACHFRAME,
};

struct UNWIND_CODE {
    u8 offset_in_prolog;
    u8 unwind_operation_code    : 4;
    u8 op_info                  : 4;
};

struct UNWIND_INFO {
    u8 version               : 3;
    u8 flags                 : 5;
    u8 size_of_prolog;
    u8 count_of_unwind_codes;
    u8 frame_register        : 4;
    u8 frame_register_offset : 4;
    struct UNWIND_CODE unwind_code[1U];
    // Optional Exception handler of chained unwind info
    //
    // Exception handler
    // ULONG address_of_exception_handler;
    // language specific handler data;
    //
    // Chained unwind info
    // ULONG function_start;
    // ULONG function_end;
    // ULONG unwind_info_address;
};

enum UNWIND_OP_INFO_TYPE {
    RAX,
    RCX,
    RDX,
    RBX,
    RSP,
    RBP,
    RSI,
    RDI,
    R8,
    R9,
    R10,
    R11,
    R12,
    R13,
    R14,
    R15,
};

struct module* get_addr_module(u64 addr) {
    for (u8 index = 0U; index < 64U; index++) {
        if (((1ULL << index) & modules_list) != 0U) {
            continue;
        }

        struct module* module = &modules[index];
        u64 offset = addr - (u64)module->base_addr;
        if (0U <= offset && offset < module->nt_header.OptionalHeader.SizeOfImage) {
            return module;
        }
    }

    return NULL;
}

void print_reg(enum UNWIND_OP_INFO_TYPE op_info) {
    switch (op_info) {
        case RAX:
            print("RAX\n");
            break;
        case RCX:
            print("RCX\n");
            break;
        case RDX:
            print("RDX\n");
            break;
        case RBX:
            print("RBX\n");
            break;
        case RSP:
            print("RSP\n");
            break;
        case RBP:
            print("RBP\n");
            break;
        case RSI:
            print("RSI\n");
            break;
        case RDI:
            print("RDI\n");
            break;
        case R8:
            print("R8\n");
            break;
        case R9:
            print("R9\n");
            break;
        case R10:
            print("R10\n");
            break;
        case R11:
            print("R11\n");
            break;
        case R12:
            print("R12\n");
            break;
        case R13:
            print("R13\n");
            break;
        case R14:
            print("R14\n");
            break;
        case R15:
            print("R15\n");
            break;
    }
}

void walk_stack(void* ctx) {
    u64 addr = ((CONTEXT*)ctx)->Rip;
    u64 rsp = ((CONTEXT*)ctx)->Rsp;
    u64 rbp = ((CONTEXT*)ctx)->Rbp;
    UNUSED_VAR(rbp); // Not used
    struct module* module = NULL;
    struct UNWIND_INFO* u_info = VirtualAlloc(NULL, 4096U, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
    while (addr) {
        module = get_addr_module(addr);

        if (module == NULL) {
            print("Module not found!\n");
            return;
        }
        // print("Module name %s\n", module->name);
        // print("Module address %xu\n", (u64)module->base_addr);
        // print("Module size %xu\n", module->nt_header.OptionalHeader.SizeOfImage);
        // print("Module code address %xu\n", (u64)module->base_addr + module->nt_header.OptionalHeader.BaseOfCode);
        // print("Module code size %xu\n", module->nt_header.OptionalHeader.SizeOfCode);
        // print("Function address %xu\n", addr);
        // print("Function offset in module %xu\n", addr - (u64)module->base_addr);

#ifdef HASHMAP
        struct element* sym = find_function(module, (u32)(addr - (u64)module->base_addr));
        if (sym == NULL) {
#else
        u32 sym_index = find_function(module, addr - (u64)module->base_addr);
        if (sym_index == (u32)-1) {
#endif
            print("No symbol found at address %x\n", addr);

            addr = *(u64*)rsp;
            continue;
        }
#ifdef HASHMAP
        void* name = (void*)sym->name;
        u64 unwind_offset = sym->unwind_offset;
        u32 start = sym->start;

        u8 buffer[4096U];
        read_memory(name, buffer, 4096U);
        print("Buffer is %s!%s\n", module->name, buffer);
#else
        u8* name = module->functions_name[sym_index];
        u64 unwind_offset = module->functions_unwind[sym_index];
        u32 start = module->functions_start[sym_index];
#endif
        print("%s!%s\n", module->base_name, name);
        // print("rsp value is %xu\n", rsp);
        // print("rbp value is %xu\n", rbp);

        if (!read_memory((const void*)((u64)module->base_addr + unwind_offset), (void*)u_info, 4096U)) {
            print("Cannot read unwind info in for function %s\n", name);
            continue;
        }

        size_t total_size = sizeof(*u_info) + (u_info->count_of_unwind_codes - 1U) * sizeof(struct UNWIND_CODE);
        assert(total_size < 4096U);

        u64 func_start = (u64)module->base_addr + start;
        u32 code_index = 0U;
        // Undo prolog
        if (func_start <= addr && addr < func_start + u_info->size_of_prolog) {
            u32 RIP_offset = addr - func_start;
            // print("RIP is in prolog\n");
            struct UNWIND_CODE* cur_code = &u_info->unwind_code[code_index];
            while (cur_code->offset_in_prolog != RIP_offset){
                cur_code = &u_info->unwind_code[++code_index];
            } 
        }

        for (; code_index < u_info->count_of_unwind_codes;) {
            struct UNWIND_CODE* cur_code = &u_info->unwind_code[code_index];
            switch (cur_code->unwind_operation_code) {
                case UWOP_PUSH_NONVOL:
                    // print("push non volatile\n");
                    // print("push ");
                    // print_reg(cur_code->op_info);
                    rsp += 8U;
                    code_index++;
                    break;
                case UWOP_ALLOC_LARGE:
                    // print("alloc large\n");
                    if (cur_code->op_info == 0U) {
                        u16 size = *(u16*)&u_info->unwind_code[code_index + 1U];
                        if (u_info->frame_register) {
                            // print("sub rbp %xu\n", size);
                            rbp += size;
                        } else {
                            // print("sub rsp %xu\n", size);
                            rsp += size;
                        }
                        code_index += 2U;
                    } else {
                        u32 size = *(u32*)&u_info->unwind_code[code_index + 1U];
                        if (u_info->frame_register) {
                            // print("sub rbp %xu\n", size);
                            rbp += size;
                        } else {
                            // print("sub rsp %xu\n", size);
                            rsp += size;
                        }
                        code_index += 3U;
                    }
                    break;
                case UWOP_ALLOC_SMALL:
                    // print("alloc small\n");
                    u32 size = cur_code->op_info * 8U + 8U;
                    if (u_info->frame_register) {
                        // print("sub rbp %xu\n", size);
                        rbp += size;
                    } else {
                        // print("sub rsp %xu\n", size);
                        rsp += size;
                    }
                    code_index++;
                    break;
                case UWOP_SET_FPREG:
                    // print("set frame point reg\n");
                    // print("mov rbp rsp - %xu\n", u_info->frame_register_offset * 16U);
                    code_index++;
                    break;
                case UWOP_SAVE_NONVOL:
                    // print("save non volatile\n");
                    if (u_info->frame_register) {
                        // print("mov rbp ");
                    } else {
                        // print("mov rsp ");
                    }
                    // print_reg(cur_code->op_info);
                    code_index += 2U;
                    break;
                case UWOP_SAVE_NONVOL_FAR:
                    // print("save non volatile far\n");
                    code_index += 3U;
                    break;
                case UWOP_SAVE_XMM128:
                    // print("save xmm128\n");
                    code_index += 2U;
                    break;
                case UWOP_SAVE_XMM128_FAR:
                    // print("save xmm128 far\n");
                    code_index += 3U;
                    break;
                case UWOP_PUSH_MACHFRAME:
                    // print("push mach frame\n");
                    code_index++;
                    break;
                default:
                    // print("bad operation code\n");
                    code_index++;
            }
        }
        u64 val;
        read_memory((void*)rsp, &val, sizeof(u64));
        rsp += 8U;
        // print("called address is %xu\n", val);
        addr = val;
    }
    VirtualFree(u_info, 4096U, MEM_RELEASE);
}