#include <Windows.h>
#include <cstdint>
#include "../third_party/mhook/mhook-lib/mhook.h"
#include "typedefs.h"
#include "abstracthook.h"
#include "mhook.h"

#pragma comment(lib, "..\\x64\\debug\\test_cases.lib")

static TypeSmall trueSmall = &_small;
static TypeBranch trueBranch = &_branch;
static TypeRip_relative trueRip_Relative = &_rip_relative;
static TypeAVX trueAVX = &_AVX;
static TypeRDRAND trueRDRAND = &_RDRAND;

AbstractHookEngine* g_mhook = new MHook();

uint64_t MHook_Hooks::hookSmall(void) {
	g_mhook->small_ = true;

	return trueSmall();
}

uint64_t MHook_Hooks::hookBranch(uint64_t x) {
	g_mhook->branch = true;

	return trueBranch(x);
}

uint64_t MHook_Hooks::hookRip_relative(void) {
	g_mhook->rip_relative = true;

	return trueRip_Relative();
}

void MHook_Hooks::_AVX(float num, void* res) {
	g_mhook->avx = true;

	return trueAVX(num, res);
}

uint32_t MHook_Hooks::_RDRAND(void) {
	g_mhook->rdrand = true;

	return trueRDRAND();
}

bool MHook::hook_all(void) {
	bool ret = Mhook_SetHook((PVOID*)&trueSmall, &MHook_Hooks::hookSmall);
	ret |= Mhook_SetHook((PVOID*)&trueBranch, &MHook_Hooks::hookBranch);
	ret |= Mhook_SetHook((PVOID*)&trueRip_Relative, &MHook_Hooks::hookRip_relative);

	ret |= Mhook_SetHook((PVOID*)&trueAVX, &MHook_Hooks::_AVX);
	ret |= Mhook_SetHook((PVOID*)&trueRDRAND, &MHook_Hooks::_RDRAND);

	return ret;
}

bool MHook::unhook_all() {
	return Mhook_Unhook((PVOID*)&trueSmall) &&
		Mhook_Unhook((PVOID*)&trueBranch) &&
		Mhook_Unhook((PVOID*)&trueAVX) &&
		Mhook_Unhook((PVOID*)&trueRDRAND);
}