#include "hooks.h"
#include "crt.h"

namespace hooks
{
	HINTERNET win_http_open(
		LPCWSTR agent,
		DWORD   access_type,
		LPCWSTR proxy,
		LPCWSTR proxy_bypass,
		DWORD   flags
	)
	{
        MessageBoxA(NULL, "WinHttpOpen Called", "INFO", NULL);
        DWORD bytes_written;
        ::WriteFile(
            winhttp_log_handle,
            L"user-agent: ",
            ::wcslen(L"user-agent: "),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            agent,
            ::wcslen(agent),
            &bytes_written, 
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            L"\n\n",
            ::wcslen(L"\n\n"),
            &bytes_written,
            NULL
        );

		return reinterpret_cast<decltype(&win_http_open)>(GetProcAddress(LoadLibraryA("Winhttp.dll"), "WinHttpOpen"))
			(
				agent, 
				access_type, 
				proxy, 
				proxy_bypass,
				flags
			);
	}

    HINTERNET win_http_open_req(
        IN HINTERNET conn,
        IN LPCWSTR   verb,
        IN LPCWSTR   obj_name,
        IN LPCWSTR   version,
        IN LPCWSTR   referrer,
        IN LPCWSTR* accepted_types,
        IN DWORD     flags
    )
    {
        MessageBoxA(NULL, "WinHttpOpenRequest Called", "INFO", NULL);
        DWORD bytes_written;
        ::WriteFile(
            winhttp_log_handle,
            L"downloading: ",
            ::wcslen(L"downloading: "),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            obj_name,
            ::wcslen(obj_name),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            L"\n\n",
            ::wcslen(L"\n\n"),
            &bytes_written,
            NULL
        );

        return reinterpret_cast<decltype(&win_http_open_req)>(GetProcAddress(LoadLibraryA("winhttp.dll"), "WinHttpOpenRequest"))
            (
                conn,
                verb,
                obj_name,
                version,
                referrer,
                accepted_types,
                flags
            );
    }

    HINTERNET win_http_connect(
        IN HINTERNET     session,
        IN LPCWSTR       server_name,
        IN INTERNET_PORT server_port,
        IN DWORD         reserve
    )
    {
        MessageBoxA(NULL, "WinHttpConnect Called", "INFO", NULL);
        DWORD bytes_written;
        ::WriteFile(
            winhttp_log_handle,
            L"connecting: ",
            ::wcslen(L"connecting: "),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            server_name,
            ::wcslen(server_name),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            L"\n\n",
            ::wcslen(L"\n\n"),
            &bytes_written,
            NULL
        );

        return reinterpret_cast<decltype(&win_http_connect)>(GetProcAddress(LoadLibraryA("winhttp.dll"), "WinHttpConnect"))
            (session, server_name, server_port, reserve);
    }

    HMODULE get_module_handle(
        LPCSTR module_name
    )
    {
        MessageBoxA(NULL, module_name, "Getting module handle of this module", NULL);
        DWORD bytes_written;
        ::WriteFile(
            winhttp_log_handle,
            L"getting module: ",
            ::wcslen(L"getting module: "),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            module_name,
            ::strlen(module_name),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            L"\n\n",
            ::wcslen(L"\n\n"),
            &bytes_written,
            NULL
        );
        return GetModuleHandleA(module_name);
    }

    FARPROC get_proc_addr(
        HMODULE module_handle,
        LPCSTR  proc_name
    )
    {
        MessageBoxA(NULL, proc_name, "Getting address of this proc", NULL);
        DWORD bytes_written;
        ::WriteFile(
            winhttp_log_handle,
            L"getting proc name: ",
            ::wcslen(L"getting proc name: "),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            proc_name,
            ::strlen(proc_name),
            &bytes_written,
            NULL
        );

        ::WriteFile(
            winhttp_log_handle,
            L"\n\n",
            ::wcslen(L"\n\n"),
            &bytes_written,
            NULL
        );

        return GetProcAddress(module_handle, proc_name);
    }

    void* iat_hook(void* base_addr, const char* import, void* func_addr)
    {
        if (!base_addr || *(short*)base_addr != 0x5A4D || !import || !func_addr)
            return NULL;

        PIMAGE_DOS_HEADER dos_headers =
            reinterpret_cast<PIMAGE_DOS_HEADER>(base_addr);

        PIMAGE_NT_HEADERS nt_headers =
            reinterpret_cast<PIMAGE_NT_HEADERS>(
                reinterpret_cast<DWORD_PTR>(base_addr) + dos_headers->e_lfanew);

        IMAGE_DATA_DIRECTORY import_dir =
            nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        PIMAGE_IMPORT_DESCRIPTOR import_des =
            reinterpret_cast<PIMAGE_IMPORT_DESCRIPTOR>(import_dir.VirtualAddress + (DWORD_PTR)base_addr);

        LPCSTR lib_name = NULL;
        PVOID result = NULL;
        PIMAGE_IMPORT_BY_NAME func_name = NULL;

        if (!import_des)
            return NULL;

        while (import_des->Name != NULL)
        {
            lib_name = (LPCSTR)import_des->Name + (DWORD_PTR)base_addr;

            if (GetModuleHandleA(lib_name))
            {
                PIMAGE_THUNK_DATA org_first_thunk = NULL, first_thunk = NULL;
                org_first_thunk = (PIMAGE_THUNK_DATA)((DWORD_PTR)base_addr + import_des->OriginalFirstThunk);
                first_thunk = (PIMAGE_THUNK_DATA)((DWORD_PTR)base_addr + import_des->FirstThunk);
                while (org_first_thunk->u1.AddressOfData != NULL)
                {
                    func_name = (PIMAGE_IMPORT_BY_NAME)((DWORD_PTR)base_addr + org_first_thunk->u1.AddressOfData);
                    if (strcmp(func_name->Name, import) == 0)
                    {
                        // save old function pointer
                        result = reinterpret_cast<PVOID>(first_thunk->u1.Function);

                        DWORD old_protect;
                        VirtualProtect(&first_thunk->u1.Function, 8, PAGE_READWRITE, &old_protect);
                        // swap address
                        first_thunk->u1.Function = reinterpret_cast<ULONG64>(func_addr);

                        VirtualProtect(&first_thunk->u1.Function, 8, old_protect, &old_protect);
                        return result;
                    }
                    ++org_first_thunk;
                    ++first_thunk;
                }
            }
            ++import_des;
        }
        return NULL;
    }
}