#include "parser.h"

#include "utils.h"
#include "module.h"
#include "stack_walk.h"
#include "disasm.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 THREAD_TERMINATE                 (0x0001)  
#define THREAD_SUSPEND_RESUME            (0x0002)  
#define THREAD_GET_CONTEXT               (0x0008)  
#define THREAD_SET_CONTEXT               (0x0010)  
#define THREAD_QUERY_INFORMATION         (0x0040)  
#define THREAD_SET_INFORMATION           (0x0020)  
#define THREAD_SET_THREAD_TOKEN          (0x0080)
#define THREAD_IMPERSONATE               (0x0100)
#define THREAD_DIRECT_IMPERSONATION      (0x0200)
#define THREAD_SET_LIMITED_INFORMATION   (0x0400)  // winnt
#define THREAD_QUERY_LIMITED_INFORMATION (0x0800)  // winnt
#define THREAD_RESUME                    (0x1000)  // winnt

//
// Debug APIs
//
#define EXCEPTION_DEBUG_EVENT       1
#define CREATE_THREAD_DEBUG_EVENT   2
#define CREATE_PROCESS_DEBUG_EVENT  3
#define EXIT_THREAD_DEBUG_EVENT     4
#define EXIT_PROCESS_DEBUG_EVENT    5
#define LOAD_DLL_DEBUG_EVENT        6
#define UNLOAD_DLL_DEBUG_EVENT      7
#define OUTPUT_DEBUG_STRING_EVENT   8
#define RIP_EVENT                   9

typedef u32 (*PTHREAD_START_ROUTINE)(void* lpThreadParameter);
typedef PTHREAD_START_ROUTINE LPTHREAD_START_ROUTINE;

#define EXCEPTION_MAXIMUM_PARAMETERS 15 // maximum number of exception parameters

typedef PCONTEXT LPCONTEXT;

#define SORT_DEFAULT                     0x0     // sorting default

#define LANG_NEUTRAL                     0x00
#define SUBLANG_SYS_DEFAULT              0x02    // system default

#define MAKELANGID(p, s)       ((((u16  )(s)) << 10) | (u16  )(p))
#define LANG_SYSTEM_DEFAULT    (MAKELANGID(LANG_NEUTRAL, SUBLANG_SYS_DEFAULT))

#define MAKELCID(lgid, srtid)  ((u32)((((u32)((u16  )(srtid))) << 16) |  \
                                         ((u32)((u16  )(lgid)))))
#define LOCALE_SYSTEM_DEFAULT  (MAKELCID(LANG_SYSTEM_DEFAULT, SORT_DEFAULT))

//
//  Compare String Return Values.
//
#define CSTR_LESS_THAN            1           // string 1 less than string 2
#define CSTR_EQUAL                2           // string 1 equal to string 2
#define CSTR_GREATER_THAN         3           // string 1 greater than string 2

#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 ======
// ========================

extern u8 expect_single_step;
extern u32 continue_status;

enum TOKEN_TYPE {
    // Commands
    TOKEN_CONTINUE,
    TOKEN_CALLSTACK,
    TOKEN_LISTMODULE,
    TOKEN_DUMP_SYMBOLS,
    TOKEN_PRINT_REG,
    TOKEN_STEP_INTO,
    TOKEN_DUMP_MEMORY,
    TOKEN_UNASSEMBLE,
    TOKEN_QUIT,

    // Primaries
    TOKEN_DECIMAL_NUMBER,
    TOKEN_HEXADECIMAL_NUMBER,

    TOKEN_UNKNOWN,
};

struct token {
    enum TOKEN_TYPE type;
    const char* lexeme;
    u8 size;
};

struct command {
    const char* start;
    u8 size;
};

struct ast_node {
    struct token* token;
    struct ast_node* left;
    struct ast_node* right;
};

#define DECLARE_COMMAND(x) { .start = x, .size = sizeof(x) }

static struct command commands[] = {
    DECLARE_COMMAND("g"),
    DECLARE_COMMAND("k"),
    DECLARE_COMMAND("lm"),
    DECLARE_COMMAND("x"),
    DECLARE_COMMAND("r"),
    DECLARE_COMMAND("t"),
    DECLARE_COMMAND("db"),
    DECLARE_COMMAND("u"),
    DECLARE_COMMAND("q"),
};

_Static_assert((sizeof(commands) / sizeof(commands[0])) == TOKEN_QUIT + 1, "Missing commands");

void lexer(const char* str_cmd, size_t str_cmd_len, struct token* tokens) {
    const char* cur = str_cmd;
    const char* end = str_cmd + str_cmd_len;
    struct token* cur_token = tokens;
    str_cmd_len = 0;

    while(cur < end) {
        cur_token->type = TOKEN_UNKNOWN;
        if(*cur == '0' && cur[1] == 'x') {
            const char* start = cur;
            cur += 2;

            while(is_hexadecimal_digit(*cur) && cur < end) {
                cur++;
            }

            cur_token->lexeme = start;
            cur_token->type = TOKEN_HEXADECIMAL_NUMBER;
            cur_token->size = (u8)(cur - start);
            cur_token++;
        } else if(is_decimal_digit(*cur)) {
            const char* start = cur;
            cur++;

            while(is_decimal_digit(*cur) && cur < end) {
                cur++;
            }

            cur_token->lexeme = start;
            cur_token->type = TOKEN_DECIMAL_NUMBER;
            cur_token->size = (u8)(cur - start);
            cur_token++;
        } else if(is_alpha(*cur)) {
            const char* start = cur;
            cur++;
            while(is_alpha(*cur) && cur < end) {
                cur++;
            }

            cur_token->lexeme = start;
            cur_token->size = (u8)(cur - start);

            for(int i = 0; i < sizeof(commands) / sizeof(commands[0]); i++) {
                struct command cur_cmd = commands[i];
                if (CompareStringA(LOCALE_SYSTEM_DEFAULT, 0, (i8*)cur_cmd.start, cur_cmd.size, (i8*)cur_token->lexeme, cur_token->size) == CSTR_EQUAL) {
                    cur_token->type = i;
                    cur_token++;
                    break;
                }
            }
        } else {
            cur++;
        }
    }
}

struct ast_node* parse_primary(struct token* tokens, struct ast_node* nodes) {
    struct ast_node* cur_node = &nodes[0];
    struct token* cur_token = &tokens[0];
    if (cur_token->type == TOKEN_HEXADECIMAL_NUMBER ||
        cur_token->type == TOKEN_DECIMAL_NUMBER) {
        cur_node->token = cur_token;
        cur_node->left = NULL;
        cur_node->right = NULL;
        return cur_node;
    } else {
        print("Expecting a primary token\n");
        return NULL;
    }
}

struct ast_node* parse_command(struct token* tokens, struct ast_node* nodes) {
    struct ast_node* cur_node = &nodes[0];
    struct token* cur_token = &tokens[0];
    switch(cur_token->type) {
        case TOKEN_CONTINUE:
            tokens++;
            print("TOKEN_CONTINUE\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_CALLSTACK:
            tokens++;
            print("TOKEN_CALLSTACK\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_LISTMODULE:
            tokens++;
            print("TOKEN_LISTMODULE\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_DUMP_SYMBOLS:
            tokens++;
            print("TOKEN_DUMP_SYMBOLS\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_PRINT_REG:
            tokens++;
            print("TOKEN_PRINT_REG\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_STEP_INTO:
            tokens++;
            print("TOKEN_STEP_INTO\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_DUMP_MEMORY:
            tokens++;
            print("TOKEN_DUMP_MEMORY\n");
            cur_node->token = cur_token;
            cur_node->left = parse_primary(tokens++, ++nodes);
            cur_node->right = parse_primary(tokens++, ++nodes);
            return cur_node;
        case TOKEN_UNASSEMBLE:
            tokens++;
            print("TOKEN_UNASSEMBLE\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        case TOKEN_QUIT:
            tokens++;
            print("TOKEN_QUIT\n");
            cur_node->token = cur_token;
            cur_node->left = NULL;
            cur_node->right = NULL;
            return cur_node;
        default:
            print("%s must be a command\n", cur_token->lexeme);
            return NULL;
    }
}

u64 eval_number(const struct ast_node* node) {
    if (node->token->type == TOKEN_DECIMAL_NUMBER ||
        node->token->type == TOKEN_HEXADECIMAL_NUMBER) {
        return atou64(node->token->lexeme, node->token->size);
    }
    return ~0U;
}

void run(struct ast_node* root) {
    extern DEBUG_EVENT dbg_event;
    extern u8 process_commands;
    extern u8 is_open;

    struct ast_node cur_node = root[0];
    CONTEXT ctx = {
        .ContextFlags = CONTEXT_ALL,
    };
    void* h_thread = OpenThread(
        THREAD_GET_CONTEXT | THREAD_SET_CONTEXT,
        FALSE,
        dbg_event.dwThreadId
    );
    GetThreadContext(h_thread, &ctx);

    if (cur_node.token->type == TOKEN_CONTINUE) {
        ContinueDebugEvent(
            dbg_event.dwProcessId,
            dbg_event.dwThreadId,
            continue_status
        );
        process_commands = FALSE;
    } else if (cur_node.token->type == TOKEN_CALLSTACK) {
        walk_stack(&ctx);
    } else if (cur_node.token->type == TOKEN_LISTMODULE) {
        for (u8 index = 0U; index < 64U; index++) {
            if (((1ULL << index) & modules_list) == 0U) {
                struct module* module = &modules[index];
                print("Module Name : %s\n", module->base_name);
                print("Module Base : %x\n", module->base_addr);
                print("Module Size : %x\n\n", module->nt_header.OptionalHeader.SizeOfImage);
            }
        }
    } else if (cur_node.token->type == TOKEN_PRINT_REG) {
        print("RIP : %.12x\n", ctx.Rip);
        print("RCX : %.12x\n", ctx.Rcx);
        print("RDX : %.12x\n", ctx.Rdx);
        print("RBX : %.12x\n", ctx.Rbx);
        print("RSP : %.12x\n", ctx.Rsp);
        print("RBP : %.12x\n", ctx.Rbp);
        print("RSI : %.12x\n", ctx.Rsi);
        print("RDI : %.12x\n", ctx.Rdi);
        process_commands = TRUE;
    } else if (cur_node.token->type == TOKEN_STEP_INTO) {
        ctx.EFlags |= 0x100; // TRAP_FLAG
        if(SetThreadContext(h_thread, &ctx) == 0) {
            __debugbreak();
        }
        ContinueDebugEvent(
            dbg_event.dwProcessId,
            dbg_event.dwThreadId,
            continue_status
        );
        expect_single_step = TRUE;
        process_commands = FALSE;
    } else if (cur_node.token->type == TOKEN_QUIT) {
        is_open = FALSE;
        process_commands = FALSE;
    } else if (cur_node.token->type == TOKEN_DUMP_MEMORY) {
        u64 addr = eval_number(cur_node.left);
        u64 size = eval_number(cur_node.right);
        u8* buf = VirtualAlloc(NULL, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);

        if(read_memory((const void*)addr, (void*)buf, size)) {
            for (u32 index = 0U; index < size; index += 8U) {
                print("%.2x %.2x %.2x %.2x %.2x %.2x %.2x %.2x\n", buf[index], buf[index + 1U], buf[index + 2U], buf[index + 3U], buf[index + 4U], buf[index + 5U], buf[index + 6U], buf[index + 7U]);
            }
            print("\n");
        } else {
            
            print("Unable to read memory at this address, error %u\n", GetLastError());
        }

        VirtualFree(buf, size, MEM_RELEASE);
    } else if (cur_node.token->type == TOKEN_DUMP_SYMBOLS) {
        for (u8 index = 0U; index < 64U; index++) {
            if (((1ULL << index) & modules_list) == 0U) {
                struct module* module = &modules[index];
#ifdef HASHMAP
                iterate(&module->symbols);
#else
                for (u32 sym_index = 0U; sym_index < module->functions_count; sym_index++) {
                    print("%s\n", module->functions_name[index]);
                }
#endif
            }
        }
    } else if (cur_node.token->type == TOKEN_UNASSEMBLE) {
        struct module* module = get_addr_module(ctx.Rip);
        u32 sym_index = find_function(module, ctx.Rip - (u64)module->base_addr);
        if (sym_index != (u32)-1) {
            u32 start = module->functions_start[sym_index];
            u32 end = module->functions_end[sym_index];
            unassemble((u8*)module->base_addr + start, end - start);
        } else {
            print("No symbol found at address %x\n", ctx.Rip);
        }
    } else {
        print("Command unknown\n");
        process_commands = TRUE;
    }

    CloseHandle(h_thread);
}

u8 interpreter(const char* str_cmd, size_t str_cmd_len) {
    struct token tokens[64U];
    struct ast_node nodes[64U];

    lexer(str_cmd, str_cmd_len, tokens);

    struct ast_node* root = parse_command(tokens, nodes);

    if(root == NULL) {
        print("Ill formed command\n");
        return FALSE;
    }

    run(root);
    return TRUE;
}