// Uses BlHooks for some things
// Todo: Make independent

#include <Windows.h>
#include <Psapi.h>

#include "lua.hpp"

// Custom types
void bll_hk_newMetatable(lua_State* L, const char* name, struct luaL_Reg* methods) {
	if(luaL_newmetatable(L, name)) {
		if(methods!=NULL)
			luaL_setfuncs(L, methods, 0);
		//lua_pushliteral(L, "__index");
		//lua_settable(L, -3);
		//lua_pushliteral(L, "__metatable");
		//lua_pushliteral(L, "access to this metatable is disabled");
		//lua_settable(L, -3);
	}
	lua_pop(L, 1);
}

// Image - Reference to a loaded binary, contains the base pointer and size
typedef struct {
	void* base;
	size_t size;
} bll_hk_image;
bll_hk_image* bll_hk_newImage(lua_State* L, void* base, size_t size) {
	bll_hk_image* img = (bll_hk_image*)lua_newuserdata(L, sizeof(bll_hk_image));
	luaL_setmetatable(L, "_bllua_hk.image");
	img->base = base;
	img->size = size;
	return img;
}
int bll_hk_image_tostring(lua_State* L) {
	bll_hk_image* img = (bll_hk_image*)luaL_checkudata(L, 1, "_bllua_hk.image");
	char buf[100];
	sprintf(buf, "image:{addr=%08X,size=%08X}",
		(unsigned int)img->base,
		(unsigned int)img->size);
	lua_pushstring(L, buf);
	return 1;
}
struct luaL_Reg bll_hk_image_methods[] = {
	{"__tostring", bll_hk_image_tostring},
	NULL,
};
void bll_hk_initTypeImage(lua_State* L) {
	bll_hk_newMetatable(L, "_bllua_hk.image", bll_hk_image_methods);
}

// Addr - Just a pointer, can be added/subtracted
// Just using numbers for now
//typedef struct {
//	void* a;
//} bll_hk_addr;
//int bll_hk_newAddr(lua_State* L, void* a) {
//	bll_hk_addr* addr = (bll_hk_addr*)lua_newuserdata(L, sizeof(bll_hk_addr));
//	luaL_setmetatable(L, "_bllua_hk.addr");
//	addr->a = a;
//	return 1;
//}
//int bll_hk_addr_tostring(lua_State* L) {
//	bll_hk_addr* addr = (bll_hk_addr*)luaL_checkudata(L, 1, "_bllua_hk.addr");
//	
//	char buf[100];
//	sprintf(buf, "addr:%08X", (unsigned int)addr->a);
//	lua_pushstring(L, buf);
//	return 1;
//}
//int bll_hk_addr_add(lua_State* L) {
//	bll_hk_addr* addr = (bll_hk_addr*)luaL_checkudata(L, 1, "_bllua_hk.addr");
//	size_t ofs = luaL_checkinteger(L, 2);
//	
//	bll_hk_newAddr(L, (char*)addr->a + ofs);
//	return 1;
//}
//int bll_hk_addr_sub(lua_State* L) {
//	bll_hk_addr* addr = (bll_hk_addr*)luaL_checkudata(L, 1, "_bllua_hk.addr");
//	size_t ofs = luaL_checkinteger(L, 2);
//	
//	bll_hk_newAddr(L, (char*)addr->a - ofs);
//	return 1;
//}
//struct luaL_Reg bll_hk_addr_methods[] = {
//	{"__tostring", bll_hk_addr_tostring},
//	{"__add", bll_hk_addr_add},
//	{"__sub", bll_hk_addr_sub},
//	NULL,
//};
//void bll_hk_initTypeAddr(lua_State* L) {
//	bll_hk_newMetatable(L, "_bllua_hk.addr", bll_hk_addr_methods);
//}


// Scanning
bool bll_hk_patternMatch(char* addr, const char* pat, const char* mask) {
	for(; *mask; ++addr, ++pat, ++mask) {
		if(*mask=='x' && *addr!=*pat)
			return false;
	}
	return !*mask;
}
void* bll_hk_scan(void* imgBase, size_t imgSize, const char* pat, const char* mask,
	int skip) {
	char* scanStart = (char*)imgBase;
	char* scanEnd = scanStart + imgSize;
	for(char* addr=scanStart; addr<scanEnd; addr++) {
		if(bll_hk_patternMatch(addr, pat, mask)) {
			if(skip==0)
				return addr;
			else
				skip--;
		}
	}
	return NULL;
}


// Lua lib functions

#define BLL_HK_PUSHADDR(_l, _a) lua_pushnumber((_l), (lua_Number)(unsigned int)(_a))
#define BLL_HK_CHECKADDR(_l, _i) (void*)luaL_checkinteger((_l), (_i))
#define BLL_HK_CHECKSIZET(_l, _i) (size_t)luaL_checkinteger((_l), (_i))

int bll_libhk_open(lua_State* L) {
	HMODULE mod;
	if(lua_isnoneornil(L, 1)) {
		mod = GetModuleHandle(NULL);
	} else {
		const char* modName = luaL_checkstring(L, 1);
		mod = GetModuleHandle(modName);
	}
	if(mod) {
		MODULEINFO info;
		GetModuleInformation(GetCurrentProcess(), mod, &info, sizeof(MODULEINFO));
		bll_hk_newImage(L, info.lpBaseOfDll, info.SizeOfImage);
	} else {
		lua_pushnil(L);
	}
	return 1;
}
int bll_libhk_scan(lua_State* L) {
	bll_hk_image* img = (bll_hk_image*)luaL_checkudata(L, 1, "_bllua_hk.image");
	const char* pat  = luaL_checklstring(L, 2, NULL);
	const char* mask = luaL_checklstring(L, 3, NULL);
	int skip = luaL_checkinteger(L, 4);
	void* addr = bll_hk_scan(img->base, img->size, pat, mask, skip);
	if(addr==NULL)
		lua_pushnil(L);
	else
		BLL_HK_PUSHADDR(L, addr);
	return 1;
}
int bll_libhk_readStr(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	size_t len = BLL_HK_CHECKSIZET(L, 2);
	lua_pushlstring(L, (const char*)addr, len); // based lua handling memory for us
	return 1;
}
int bll_libhk_writeStr(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	size_t len = 0;
	const char* data = luaL_checklstring(L, 2, &len);
	memcpy(addr, data, len);
	BLL_HK_PUSHADDR(L, (char*)addr+len);
	return 1;
}
int bll_libhk_readChar(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	char data = *(char*)addr;
	lua_pushinteger(L, data);
	return 1;
}
int bll_libhk_writeChar(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	char data = (char)luaL_checkinteger(L, 2);
	*(char*)addr = data;
	BLL_HK_PUSHADDR(L, (char*)addr + 1);
	return 1;
}
int bll_libhk_readShort(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	short data = *(short*)addr;
	lua_pushinteger(L, data);
	return 1;
}
int bll_libhk_writeShort(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	short data = (short)luaL_checkinteger(L, 2);
	*(short*)addr = data;
	BLL_HK_PUSHADDR(L, (char*)addr + 2);
	return 1;
}
int bll_libhk_readInt(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	int data = *(int*)addr;
	lua_pushinteger(L, data);
	return 1;
}
int bll_libhk_writeInt(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	int data = (int)luaL_checkinteger(L, 2);
	*(int*)addr = data;
	BLL_HK_PUSHADDR(L, (char*)addr + 4);
	return 1;
}
int bll_libhk_readFloat(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	float data = *(float*)addr;
	lua_pushnumber(L, data);
	return 1;
}
int bll_libhk_writeFloat(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	float data = luaL_checknumber(L, 2);
	*(float*)addr = data;
	BLL_HK_PUSHADDR(L, (char*)addr + 4);
	return 1;
}
int bll_libhk_readDouble(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	double data = *(double*)addr;
	lua_pushnumber(L, data);
	return 1;
}
int bll_libhk_writeDouble(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	double data = luaL_checknumber(L, 2);
	*(double *)addr = data;
	BLL_HK_PUSHADDR(L, (char*)addr + 8);
	return 1;
}
int bll_libhk_protect(lua_State* L) {
	void* addr = BLL_HK_CHECKADDR(L, 1);
	size_t len = BLL_HK_CHECKSIZET(L, 2);
	DWORD newProt = (DWORD)luaL_checkinteger(L, 3);
	DWORD oldProt = 0;
	VirtualProtect(addr, len, newProt, &oldProt);
	lua_pushinteger(L, (lua_Integer)oldProt);
	return 1;
}
int bll_libhk_malloc(lua_State* L) {
	size_t size = (size_t)luaL_checkinteger(L, 1);
	void* data = malloc(size);
	BLL_HK_PUSHADDR(L, data);
	return 1;
}
int bll_libhk_free(lua_State* L) {
	void* data = BLL_HK_CHECKADDR(L, 1);
	free(data);
	return 0;
}
//typedef void(*bll_func_ptr)(void);
//int bll_libhk_call(lua_State* L) {
//	void* jdest = BLL_HK_CHECKADDR(L, 1);
//	((bll_func_ptr)jdest)();
//	return 0;
//}

// Hooks lib

// regsPtr in ecx, L in edx
__declspec(fastcall) void bll_hk_callback(void* regsPtr, lua_State* L) {
	lua_getglobal(L, "_bllua_hk_callback");
	BLL_HK_PUSHADDR(L, regsPtr);
	lua_pcall(L, 1, 0, 0);
}
int bll_libhk_getLuaStatePtr(lua_State* L) {
	BLL_HK_PUSHADDR(L, (void*)L);
	return 1;
}
int bll_libhk_getCallbackPtr(lua_State* L) {
	BLL_HK_PUSHADDR(L, (void*)bll_hk_callback);
	return 1;
}
int bll_libhk_getStrPtr(lua_State* L) {
	const char* str = luaL_checkstring(L, 1);
	BLL_HK_PUSHADDR(L, (void*)str);
	return 1;
}

const luaL_Reg bll_hk_reg[] = {
	{"_openRaw"   , bll_libhk_open       },
	{"_scanRaw"   , bll_libhk_scan       },
	{"readStr"    , bll_libhk_readStr    },
	{"writeStr"   , bll_libhk_writeStr   },
	{"protect"    , bll_libhk_protect    },
	{"readChar"   , bll_libhk_readChar   },
	{"writeChar"  , bll_libhk_writeChar  },
	{"readShort"  , bll_libhk_readShort  },
	{"writeShort" , bll_libhk_writeShort },
	{"readInt"    , bll_libhk_readInt    },
	{"writeInt"   , bll_libhk_writeInt   },
	{"readFloat"  , bll_libhk_readFloat  },
	{"writeFloat" , bll_libhk_writeFloat },
	{"readDouble" , bll_libhk_readDouble },
	{"writeDouble", bll_libhk_writeDouble},
	{"malloc"     , bll_libhk_malloc     },
	{"free"       , bll_libhk_free       },
	//{"call"       , bll_libhk_call       },
	{"getStrPtr"  , bll_libhk_getStrPtr  },
	{"_getLuaStatePtr", bll_libhk_getLuaStatePtr},
	{"_getCallbackPtr", bll_libhk_getCallbackPtr},
	NULL,
};
extern "C" int __declspec(dllexport) luaopen_luahooks32core(lua_State* L) {
	luaL_register(L, "luahooks32core", bll_hk_reg);
	bll_hk_initTypeImage(L);
	//bll_hk_initTypeAddr(L);
	return 1;
}