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

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

extern uint64_t __readgsqword(unsigned long offset);

typedef struct _LIST_ENTRY {
    struct _LIST_ENTRY *Flink;
    struct _LIST_ENTRY *Blink;
} LIST_ENTRY, *PLIST_ENTRY, *PRLIST_ENTRY;

typedef struct _UNICODE_STRING {
  u16       Length;
  u16       MaximumLength;
  u16*      Buffer;
} UNICODE_STRING, *PUNICODE_STRING;

typedef struct _LDR_DATA_TABLE_ENTRY {
    void*           Reserved1[2];
    LIST_ENTRY      InMemoryOrderLinks;
    void*           Reserved2[2];
    void*           DllBase;
    void*           Reserved3[2];
    UNICODE_STRING  FullDllName;
    UNICODE_STRING  BaseDllName;
    void*           Reserved5[2];
    union
    {
        u32         CheckSum;
        void*       Reserved6;
    };
    u32             TimeDateStamp;
} LDR_DATA_TABLE_ENTRY, *PLDR_DATA_TABLE_ENTRY;

typedef struct _PEB_LDR_DATA {
    u8          Reserved1[8];
    void*       Reserved2[3];
    LIST_ENTRY  InMemoryOrderModuleList;
} PEB_LDR_DATA, *PPEB_LDR_DATA;

typedef struct _RTL_USER_PROCESS_PARAMETERS {
    u8             Reserved1[16];
    void*          Reserved2[10];
    UNICODE_STRING ImagePathName;
    UNICODE_STRING CommandLine;
} RTL_USER_PROCESS_PARAMETERS, *PRTL_USER_PROCESS_PARAMETERS;

typedef void(*PS_POST_PROCESS_INIT_ROUTINE)(void);
typedef PS_POST_PROCESS_INIT_ROUTINE* PPS_POST_PROCESS_INIT_ROUTINE;

typedef struct _PEB {
    u8 Reserved1[2];
    u8 BeingDebugged;
    u8 Reserved2[21];
    PPEB_LDR_DATA LoaderData;
    PRTL_USER_PROCESS_PARAMETERS ProcessParameters;
    u8 Reserved3[520];
    PPS_POST_PROCESS_INIT_ROUTINE PostProcessInitRoutine;
    u8 Reserved4[136];
    u32 SessionId;
} PEB, *PPEB;

typedef struct _TEB {
    void* Reserved1[12];
    PPEB  ProcessEnvironmentBlock;
    void* Reserved2[399];
    u8    Reserved3[1952];
    void* TlsSlots[64];
    u8    Reserved4[8];
    void* Reserved5[26];
    void* ReservedForOle;
    void* Reserved6[4];
    void* TlsExpansionSlots;
} TEB, *PTEB;

typedef struct _IMAGE_EXPORT_DIRECTORY {
    u32   Characteristics;
    u32   TimeDateStamp;
    u16    MajorVersion;
    u16    MinorVersion;
    u32   Name;
    u32   Base;
    u32   NumberOfFunctions;
    u32   NumberOfNames;
    u32   AddressOfFunctions;     // RVA from base of image
    u32   AddressOfNames;         // RVA from base of image
    u32   AddressOfNameOrdinals;  // RVA from base of image
} IMAGE_EXPORT_DIRECTORY, *PIMAGE_EXPORT_DIRECTORY;

//
//  Code Page Default Values.
//  Please Use Unicode, either UTF-16 (as in WCHAR) or UTF-8 (code page CP_ACP)
//
#define CP_ACP                    0           // default to ANSI code page
#define CP_OEMCP                  1           // default to OEM  code page
#define CP_MACCP                  2           // default to MAC  code page
#define CP_THREAD_ACP             3           // current thread's ANSI code page
#define CP_SYMBOL                 42          // SYMBOL translations

#define CP_UTF7                   65000       // UTF-7 translation
#define CP_UTF8                   65001       // UTF-8 translation

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

#define KERNEL32(name)      \
    PFN_##name name = NULL;

#define KERNEL_BASE(name)   \
    PFN_##name name = NULL;

#define NTDLL(name)         \
    PFN_##name name = NULL;

#include "funcs.h"

#undef NTDLL
#undef KERNEL_BASE
#undef KERNEL32

size_t str_len(const char* str) {
    const char* c = str;
    size_t count = 0U;
    while (*c != '\0') {
        c++;
        count++;
    }
    return count;
}

i32 str_cmp(const i8* str1, size_t str1_len, const i8* str2, size_t str2_len) {
    u32 min_len = str1_len < str2_len ? str1_len : str2_len;

    for (u32 index = 0U; index < min_len; index++) {
        if (str1[index] < str2[index]) {
            return -1;
        } else if (str1[index] > str2[index]) {
            return 1;
        }
    }

    return 0;
}

i32 wstr_cmp(const u16* str1, size_t str1_len, const u16* str2, size_t str2_len) {
    u32 min_len = str1_len < str2_len ? str1_len : str2_len;

    for (u32 index = 0U; index < min_len; index++) {
        if (str1[index] < str2[index]) {
            return -1;
        } else if (str1[index] > str2[index]) {
            return 1;
        }
    }

    return 0;
}

PEB* get_PEB(void) {
    return (PEB*)__readgsqword(__builtin_offsetof(TEB, ProcessEnvironmentBlock));
}

void* get_module(u16* name, size_t name_size) {
    PEB* PEB = get_PEB();
    LIST_ENTRY* first_entry = PEB->LoaderData->InMemoryOrderModuleList.Flink;
    LIST_ENTRY* cur_entry = first_entry;

    do {
        LDR_DATA_TABLE_ENTRY* table_entry = (LDR_DATA_TABLE_ENTRY*)((u8*)cur_entry - 0x10);
        UNICODE_STRING* BaseDllName =  &table_entry->BaseDllName;

        if (BaseDllName->Buffer) {
            if (wstr_cmp(name, name_size / sizeof(u16), BaseDllName->Buffer, BaseDllName->Length / sizeof(u16)) == 0) {
                return table_entry->DllBase;
            }
        }

        cur_entry = cur_entry->Flink;
    } while(cur_entry != first_entry);

    return NULL;
}

void* load_func(u64 base_addr, const char* func_name, u32 func_name_len) {
    IMAGE_DOS_HEADER* dos_header = (IMAGE_DOS_HEADER*)base_addr;
    IMAGE_NT_HEADERS64* nt_header = (IMAGE_NT_HEADERS64*)(base_addr + dos_header->e_lfanew);
    IMAGE_DATA_DIRECTORY export_table = nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
    const IMAGE_EXPORT_DIRECTORY* export_dir = (const IMAGE_EXPORT_DIRECTORY*)(base_addr + export_table.VirtualAddress);

    u16* ordinals = (u16*)(base_addr + export_dir->AddressOfNameOrdinals);
    u32* names_address = (u32*)(base_addr + export_dir->AddressOfNames);
    u32* addresses = (u32*)(base_addr + export_dir->AddressOfFunctions);

    for (u32 name_index = 0U; name_index < export_dir->NumberOfNames; name_index++) {
        const char* name = (const char*)(base_addr + names_address[name_index]);
        u32 length = str_len(name);

        if (str_cmp((i8*)name, length, (i8*)func_name, func_name_len) == 0) {
            return (void*)(base_addr + addresses[ordinals[name_index]]);
        }
    }

    return NULL;
}

void load_funcs(void) {
    u16 kernel32_str[] = L"KERNEL32.DLL";
    u16 kernel_base_str[] = L"KERNELBASE.dll";
    u16 ntdll_str[] = L"ntdll.dll";
    void* kernel32 = get_module(kernel32_str, sizeof(kernel32_str));
    void* kernel_base = get_module(kernel_base_str, sizeof(kernel_base_str));
    void* ntdll = get_module(ntdll_str, sizeof(ntdll_str));

#define KERNEL32(name)                                                    \
    name = (PFN_##name)load_func((u64)kernel32, #name, sizeof(#name));

#define KERNEL_BASE(name)                                                 \
    name = (PFN_##name)load_func((u64)kernel_base, #name, sizeof(#name));

#define NTDLL(name)                                                       \
    name = (PFN_##name)load_func((u64)ntdll, #name, sizeof(#name));

#include "funcs.h"

#undef NTDLL
#undef KERNEL_BASE
#undef KERNEL32

}