BL3Lua/BL3Lua.cpp

#include "BL3Lua.h"
#include "LuaUtility.h"

#include "include/patternscan.hpp"

int lua_exception_handler(lua_State* L, sol::optional<const std::exception&> maybe_exception, sol::string_view description) {
	// L is the lua state, which you can wrap in a state_view if necessary
	// maybe_exception will contain exception, if it exists
	// description will either be the what() of the exception or a description saying that we hit the general-case catch(...)
	std::cout << "An exception occurred in a function, here's what it says ";
	if (maybe_exception) {
		std::cout << "(straight from the exception): ";
		const std::exception& ex = *maybe_exception;
		std::cout << ex.what() << std::endl;
	}
	else {
		std::cout << "(from the description parameter): ";
		std::cout.write(description.data(), static_cast<std::streamsize>(description.size()));
		std::cout << std::endl;
	}

	// you must push 1 element onto the stack to be
	// transported through as the error object in Lua
	// note that Lua -- and 99.5% of all Lua users and libraries -- expects a string
	// so we push a single string (in our case, the description of the error)
	return sol::stack::push(L, description);
}

inline void my_panic(sol::optional<std::string> maybe_msg) {
	std::cerr << "Lua is in a panic state and will now abort() the application" << std::endl;
	if (maybe_msg) {
		const std::string& msg = maybe_msg.value();
		std::cerr << "\terror message: " << msg << std::endl;
	}
	// When this function exits, Lua will exhibit default behavior and abort()
}

BL3Lua::BL3Lua()
{
	InitializeCriticalSection(&luaCritSec);
	/*InitializeCriticalSection(&procEventCritSec);
	InitializeCriticalSection(&consoleCommandCritSec);*/
	
	lua = new sol::state();
	lua->open_libraries();
	lua->set_exception_handler(lua_exception_handler);
	lua->set_panic(sol::c_call<decltype(&my_panic), &my_panic>);

	registerBase();
	while (SDK::GEngine == nullptr)
	{
		auto engineAddress = PatternScan::FindSignature(NULL, "48 8b 4c 24 50 48 89 05 ?? ?? ?? ?? 48 85 c9 74 05 e8");
		if (engineAddress != NULL)
		{
			auto tmp = (engineAddress + 5);
			auto tmp2 = tmp + 3;
			SDK::GEngine = reinterpret_cast<SDK::UEngine**>((tmp + 7 + *(int*)(tmp2)));
		}
	}
	consoleEnvironment = sol::environment(*lua, sol::create, lua->globals());
}

BL3Lua::~BL3Lua()
{
	lua_close(lua->lua_state());
	DeleteCriticalSection(&luaCritSec);
}

bool BL3Lua::ExecuteScriptFile(const char* filePath, bool printInGame)
{
	bool result = false;
	EnterCriticalSection(&luaCritSec);
	int indexToDelete = -1;
	for (int i = 0; i < loadedComponents.size(); i++) {
		if (!strcmp(loadedComponents[i]->m_FileName.c_str(), filePath)) {
			indexToDelete = i;
			break;
		}
	}
	// TODO: Properly delete references during iteration
	// Implement iterator instead of loop
	if (indexToDelete != -1) {
		// Remove registered Hotkeys from this env
		auto hkIT = registeredHotkeys.begin();
		while (hkIT != registeredHotkeys.end()) {
			/*for (auto& r : hkIT) {*/
			sol::reference eRef = std::get<0>(*hkIT);
			sol::environment env = sol::get_environment(eRef);
			if (env == loadedComponents[indexToDelete]->m_Env) {
				eRef.abandon();
				hkIT = registeredHotkeys.erase(hkIT);
				//registeredHotkeys.erase(std::remove(registeredHotkeys.begin(), registeredHotkeys.end(), r), registeredHotkeys.end());
			}
			if(hkIT != registeredHotkeys.end())
				++hkIT;
		}
		auto procEventIT = procEventHooks.begin();
		// Remove procEventHooks from this env
		//for (auto& r : procEventHooks) {
		
		while(procEventIT != procEventHooks.end()){
			bool didErase = false;
			for (auto& p : procEventIT->second) {
				sol::reference eRef = std::get<0>(p);
				sol::environment env = sol::get_environment(eRef);
				if (env == loadedComponents[indexToDelete]->m_Env) {
					eRef.abandon();
					//r.second.erase(std::remove(r.second.begin(), r.second.end(), p), r.second.end());
					didErase = true;
				}
			}	
			if (didErase) {
				procEventIT = procEventHooks.erase(procEventIT);
				if (procEventIT != procEventHooks.end())
					++procEventIT;
			}
		}
		// Remove consoleCommands from this env
		auto ccIT = consoleCommands.begin();
		//for (auto& r : consoleCommands) {	
		while(ccIT != consoleCommands.end()){
			sol::reference eRef = std::get<0>(ccIT->second);
			sol::environment env = sol::get_environment(eRef);
			if (env == loadedComponents[indexToDelete]->m_Env) {
				eRef.abandon();
				ccIT = consoleCommands.erase(ccIT);
				//consoleCommands.erase(r.first);
			}
			if (ccIT != consoleCommands.end())
				++ccIT;
		}
		delete loadedComponents[indexToDelete];
		loadedComponents.erase(loadedComponents.begin() + indexToDelete);
	}
	lua->collect_garbage();
	LuaComponent* nComp = new LuaComponent(filePath, lua);
	nComp->Init();
	loadedComponents.push_back(nComp);
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
	return nComp->m_Env != sol::lua_nil;
}

bool BL3Lua::ExecuteScriptString(const char* scriptCode, bool printInGame)
{	
	bool result = false;
	EnterCriticalSection(&luaCritSec);
	try {
		auto func_result = lua->safe_script(scriptCode, consoleEnvironment);
		if (!func_result.valid()) {
			sol::error err = func_result;
			std::cout << "[LUA] failed with error: " << err.what() << std::endl;
			if (printInGame) {
				sol::protected_function printFunc = (*(lua))["PrintGameConsole"];
				if (printFunc.valid()) {
					std::string errorMessage = std::string("[LUA] failed with error: ").append(err.what());
					printFunc(errorMessage);
				}
			}			
		}
		else {
			result = true;
		}		
	}
	catch (...) {
	}
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
	return result;
}

void BL3Lua::RunCallbacks(const char* callbackName, sol::object data)
{
	EnterCriticalSection(&luaCritSec);
	try{
		sol::protected_function runCallbackFunc = (*lua)["RunCallbacks"];
		if (runCallbackFunc.valid()){
			sol::protected_function_result result = runCallbackFunc(callbackName, data);
			if (!result.valid()) {
				sol::error err = result;
				std::cout << "Error during callbacks: " << err.what() << std::endl;
			}
		}
	}
	catch(...){}
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::TDump(sol::object tableToDump)
{
	EnterCriticalSection(&luaCritSec);
	try {
		sol::protected_function tdumpFunc = (*lua)["internalTDump"];
		if (tdumpFunc.valid()) {
			sol::protected_function_result result = tdumpFunc(tableToDump);
			if (!result.valid()) {
				sol::error err = result;
				std::cout << "Error during tdump: " << err.what() << std::endl;
			}
		}
	}
	catch (...) {}
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::ClearCallbacks()
{
	EnterCriticalSection(&luaCritSec);
	while (callbacks && !callbacks.empty()) {
		callbacks.clear();
	}
	callbacks = lua->create_named_table("Callbacks");
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::ValidateHotkeys()
{
	EnterCriticalSection(&luaCritSec);
	std::vector<int> toDelete;
	std::vector<std::tuple<sol::reference, lua_State*, std::string, std::string, std::string>> tmp;
	bool needsFixup = false;
	for (int i = 0; i < registeredHotkeys.size(); ++i) {
		if (!std::get<0>(registeredHotkeys[i]).valid())
		{
			needsFixup = true;
		}
		else {
			tmp.push_back(registeredHotkeys[i]);
		}
	}
	if (needsFixup) {
		registeredHotkeys = tmp;
	}
	LeaveCriticalSection(&luaCritSec);
}

// Registering works
uint32_t BL3Lua::RegisterHotkey(sol::this_state s, sol::this_environment e, sol::variadic_args va) {
	bool anyError = false;
	ValidateHotkeys();
	EnterCriticalSection(&luaCritSec);
	std::vector<std::string> hotkeys;
	sol::function func;
	for (auto v : va) {
		if (v.is<std::string_view>()) {
			std::string hk = v.as<std::string>();
			if (!LuaUtility::ContainsKey(hk))
			{
				anyError = true;
				sol::protected_function printFunc = (*(lua))["PrintGameConsole"];
				if (printFunc.valid()) {
					char buffer[1024];
					sprintf_s(buffer, sizeof(buffer), "\"%s\" is not a valid keybind!", hk.c_str());
					std::string errorMessage = std::string(buffer);
					printFunc(errorMessage);
				}
			}
			else {
				hotkeys.push_back(hk);
			}
		}
		else if (v.is<sol::function>()) {
			func = v.as<sol::function>();
		}
	}
	//std::vector<std::string> hotkeys = hotkeyTable.as<std::vector<std::string>>();
	if (hotkeys.size() >= 1 && hotkeys.size() <= 3 && func.valid()) {
		for (auto& hk : hotkeys) {
			if (!LuaUtility::ContainsKey(hk))
			{
				anyError = true;
				sol::protected_function printFunc = (*(lua))["PrintGameConsole"];
				if (printFunc.valid()) {
					char buffer[1024];
					sprintf_s(buffer, sizeof(buffer), "\"%s\" is not a valid keybind!", hk.c_str());
					std::string errorMessage = std::string(buffer);
					printFunc(errorMessage);
				}
			}
		}
		if (anyError) {
			LeaveCriticalSection(&luaCritSec);
			return 0;
		}

		int index = -1;
		for (int i = 0; i < registeredHotkeys.size(); ++i) {
			sol::function testFunc = std::get<0>(registeredHotkeys[i]);
			if (std::get<0>(registeredHotkeys[i]) == func) {
				index = i;
				break;
			}
		}

		switch (hotkeys.size())
		{
		case 1:
			if (index != -1) {
				registeredHotkeys[index] = std::make_tuple(func, s.L, hotkeys.at(0), "", "");
			}
			else {
				registeredHotkeys.push_back(std::make_tuple(func, s.L, hotkeys.at(0), "", ""));
			}			
			break;
		case 2:
			if (index != -1) {
				registeredHotkeys[index] = std::make_tuple(func, s.L, hotkeys.at(0), "", "");
			}
			else {
				registeredHotkeys.push_back(std::make_tuple(func, s.L, hotkeys.at(0), hotkeys.at(1), ""));
			}
			break;
		case 3:
			if (index != -1) {
				registeredHotkeys[index] = std::make_tuple(func, s.L, hotkeys.at(0), "", "");
			}
			else {
				registeredHotkeys.push_back(std::make_tuple(func, s.L, hotkeys.at(0), hotkeys.at(1), hotkeys.at(2)));
			}
			break;
		default:
			break;
		}
	}
	LeaveCriticalSection(&luaCritSec);
	return 0;
}

void BL3Lua::RemoveHotkey(sol::this_state s, sol::this_environment e, uint32_t, sol::function func) {
	lua->collect_garbage();
}

void BL3Lua::RemoveAllHotkeys() {
	EnterCriticalSection(&luaCritSec);
	registeredHotkeys.clear();
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::AddHook(sol::this_state s, sol::this_environment e, const std::string& hookName, sol::function func, bool preHook)
{
	EnterCriticalSection(&luaCritSec);
	procEventHooks[hookName].push_back(std::make_tuple(func, s.L, preHook));
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::RemoveAllHooks()
{
	EnterCriticalSection(&luaCritSec);
	if (procEventHooks.size() > 0) {
		procEventHooks.clear();
	}
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::RemoveHook(sol::this_state s, sol::this_environment e, const std::string& hookName, sol::function func)
{
	int index = -1;
	EnterCriticalSection(&luaCritSec);
	for (int i = 0; i < procEventHooks[hookName].size(); ++i) {
		if (std::get<0>(procEventHooks[hookName][i]).registry_index() == func.registry_index()) {
			index = i;
			break;
		}			
	}
	if (index != -1)
		procEventHooks[hookName].erase(procEventHooks[hookName].begin(), procEventHooks[hookName].begin() + index);
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::AddConsoleCommand(sol::this_state s, sol::this_environment e, const std::string& command, sol::function func)
{
	EnterCriticalSection(&luaCritSec);
	consoleCommands[command] = std::make_tuple(func, s.L, command.size());
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::RemoveAllConsoleCommands()
{
	EnterCriticalSection(&luaCritSec);
	if (consoleCommands.size() > 0) {
		consoleCommands.clear();
	}
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::RemoveConsoleCommand(sol::this_state s, sol::this_environment e, const std::string& command, sol::function func)
{
	int index = -1;
	EnterCriticalSection(&luaCritSec);
	for (int i = 0; i < consoleCommands.size(); ++i) {
		if (std::get<0>(consoleCommands[command]) == func) {
			index = i;
			break;
		}
	}
	if (index != -1){
		consoleCommands.erase(command);		
	}
	lua->collect_garbage();
	LeaveCriticalSection(&luaCritSec);
}

void BL3Lua::registerBase()
{
	lua->set_function("IsKeyPressed", sol::overload(
		[](std::string keyName) { return LuaUtility::IsKeyPressed(keyName); },
		[](std::string keyName, std::string keyName2) {return LuaUtility::IsKeyPressed(keyName, keyName2); },
		[](std::string keyName, std::string keyName2, std::string keyName3) {return LuaUtility::IsKeyPressed(keyName, keyName2, keyName3);
		}));

	lua->set_function("FindObject", [](std::string name) {
		auto tmp = SDK::UObject::FindObject<SDK::UObject>(name);
		return tmp;
		});

	lua->set_function("StaticFindObject", [](SDK::UClass* pClass, const std::wstring_view name) 
		{
			return StaticFindObject(pClass, reinterpret_cast<SDK::UObject*>(-1), &name.data()[0], false);
		});

	lua->set_function("RegisterHotkey", [&](sol::this_state s, sol::this_environment e, sol::variadic_args va) {
		RegisterHotkey(s, e, va);
		});

	lua->set_function("ClearHotkeys", [&]() {RemoveAllHotkeys(); });

	sol::protected_function_result result = lua->safe_script(R"(
function RegisterCallback(onEventName, lCallback)
	while true do
		if Callbacks then
			break
		end
	end
	if Callbacks then
		Callbacks[lCallback] = onEventName
	else
		print("Callback-Table does not exist!")
	end
end

function RunCallbacks(eventName, data)
	for key, value in pairs(Callbacks) do
		key(data)
	end
end

function internalTDump(object)
	if object == nil then
		return
	end
	for key, value in pairs(getmetatable(object)) do
		if not string.find(key, "_") and key ~= "new" then
			-- if type(value) == "function" then
			-- 	print("[Function] " .. tostring(key))
			-- else
			--	print(tostring(key))
			-- end
			print(tostring(key))
		end
	end
end

defPrint = print
print = function(...)
	for i,v in ipairs{...} do
		if PrintGameConsole then
			PrintGameConsole(tostring(v))
		end
	end
	defPrint(...)
end

)");
	if (!result.valid()) {
		sol::error err = result;
		std::cout << err.what() << std::endl;
	}

	SDK::UObject::SetupLuaBind(lua);
	loadlibs();
	adjustRequirePath();
	removeAbusableFuncs();
	callbacks = lua->create_named_table("Callbacks");
}

void BL3Lua::loadlibs()
{
	const char* jsonlua =
#include "include/json.lua"
		;
	const char* splitlua =
#include "include/split.lua"
		;
	lua->require_script("json", jsonlua);
	lua->require_script("split", splitlua);
}

void BL3Lua::adjustRequirePath() {
	const std::string defaultPath = (*lua)["package"]["path"];
	wchar_t pBuffer[MAX_PATH];
	GetCurrentDirectory(sizeof(pBuffer), pBuffer);
	std::wstring widePath(pBuffer);
	std::string path(widePath.begin(), widePath.end());

	path.append("\\lua\\Modules\\?.lua");
	(*lua)["package"]["path"] = defaultPath + ";" + path;
}

void BL3Lua::removeAbusableFuncs() {
	sol::table global = lua->globals();
	if (global.valid()) {
		global["os"]["remove"] = sol::lua_nil;
	}
}