feat: lua function for getting key state

.gitignore

@@ -1,3 +1,4 @@
 .vs
 x64
 *.vcxproj.*
+vcpkg_installed

SM-KeyAPI/SM-KeyAPI.vcxproj

@@ -72,7 +72,23 @@
   <PropertyGroup Label="UserMacros" />
   <PropertyGroup Label="Vcpkg">
     <VcpkgEnableManifest>true</VcpkgEnableManifest>
-    <VcpkgEnabled>false</VcpkgEnabled>
+    <VcpkgEnabled>true</VcpkgEnabled>
+  </PropertyGroup>
+  <PropertyGroup Label="Vcpkg" Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
+    <VcpkgUseStatic>true</VcpkgUseStatic>
+    <VcpkgUseMD>true</VcpkgUseMD>
+  </PropertyGroup>
+  <PropertyGroup Label="Vcpkg" Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">
+    <VcpkgUseStatic>true</VcpkgUseStatic>
+    <VcpkgUseMD>true</VcpkgUseMD>
+  </PropertyGroup>
+  <PropertyGroup Label="Vcpkg" Condition="'$(Configuration)|$(Platform)'=='Debug|x64'">
+    <VcpkgUseStatic>true</VcpkgUseStatic>
+    <VcpkgUseMD>true</VcpkgUseMD>
+  </PropertyGroup>
+  <PropertyGroup Label="Vcpkg" Condition="'$(Configuration)|$(Platform)'=='Release|x64'">
+    <VcpkgUseStatic>true</VcpkgUseStatic>
+    <VcpkgUseMD>true</VcpkgUseMD>
   </PropertyGroup>
   <ItemDefinitionGroup Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">
     <ClCompile>
@@ -142,7 +158,9 @@
     <ClCompile Include="src\dllmain.cpp" />
   </ItemGroup>
   <ItemGroup>
-    <ClInclude Include="include\log.h" />
+    <ClInclude Include="include\sm\lua.h" />
+    <ClInclude Include="include\sm\offsets.h" />
+    <ClInclude Include="include\utils.h" />
   </ItemGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">

SM-KeyAPI/include/sm/lua.h

@@ -0,0 +1,55 @@
+#pragma once
+
+#include <string>
+
+struct lua_State; // Forward declaration, the game knows the real definition
+struct lua_CFunction; // Forward declaration, the game knows the real definition
+
+typedef void(__cdecl* lua_pushcclosure_t)(lua_State*, lua_CFunction*, int);
+typedef void(__cdecl* lua_modfield_t)(lua_State*, int, const char*);
+typedef void(__cdecl* lua_getglobal_t)(lua_State*, const char*);
+typedef void(__cdecl* lua_pushboolean_t)(lua_State*, int);
+typedef int(__cdecl* lua_modinteger_t)(lua_State*, int);
+typedef int(__cdecl* luaL_loadstring_t)(lua_State*, const char*);
+typedef int(__cdecl* lua_pcall_t)(lua_State*, int, int, int);
+typedef const char* (__cdecl* luaL_checklstring_t)(lua_State*, int, size_t*);
+typedef void(__cdecl* lua_pushinteger_t)(lua_State*, int);
+
+static constexpr int LUA_GLOBALSINDEX = -10002;
+#define lua_getglobal(L,s) lua_getfield(L, LUA_GLOBALSINDEX, s)
+
+template <typename T>
+T GetLuaFunction(const std::string& name) {
+	static HMODULE hLuaDLL = GetModuleHandleA("lua51.dll");
+	if (!hLuaDLL) {
+		std::cout << "lua51.dll not found!" << std::endl;
+		return nullptr;
+	}
+	return (T)GetProcAddress(hLuaDLL, name.c_str());
+}
+
+static auto luaL_loadstring = GetLuaFunction<luaL_loadstring_t>("luaL_loadstring");
+static auto lua_pcall = GetLuaFunction<lua_pcall_t>("lua_pcall");
+static auto lua_pushcclosure = GetLuaFunction<lua_pushcclosure_t>("lua_pushcclosure");
+static auto lua_setfield = GetLuaFunction<lua_modfield_t>("lua_setfield");
+static auto lua_getfield = GetLuaFunction<lua_modfield_t>("lua_getfield");
+static auto lua_pushboolean = GetLuaFunction<lua_pushboolean_t>("lua_pushboolean");
+static auto lua_tointeger = GetLuaFunction<lua_modinteger_t>("lua_tointeger");
+static auto lua_pushinteger = GetLuaFunction<lua_modinteger_t>("lua_pushinteger");
+static auto luaL_checklstring = GetLuaFunction<luaL_checklstring_t>("luaL_checklstring");
+
+bool hasAllLuaFunctions() {
+#define CHECK_LUA_FUNC(func) if (!func) { std::cout << #func << " not found!" << std::endl; failed = true; }
+	bool failed = false;
+	CHECK_LUA_FUNC(luaL_loadstring);
+	CHECK_LUA_FUNC(lua_pcall);
+	CHECK_LUA_FUNC(lua_pushcclosure);
+	CHECK_LUA_FUNC(lua_setfield);
+	CHECK_LUA_FUNC(lua_getfield);
+	CHECK_LUA_FUNC(lua_pushboolean);
+	CHECK_LUA_FUNC(lua_tointeger);
+	CHECK_LUA_FUNC(lua_pushinteger);
+	CHECK_LUA_FUNC(luaL_checklstring);
+#undef CHECK_LUA_FUNC
+	return !failed;
+}

SM-KeyAPI/include/sm/offsets.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <cstdint>
+
+namespace SM::Offsets {
+	constexpr uintptr_t loadLuaEnv = 0x54A7F0;
+
+	namespace Rebased {
+		const uintptr_t loadLuaEnv = (uintptr_t)GetModuleHandle(NULL) + SM::Offsets::loadLuaEnv;
+	} // namespace Rebased
+} // namespace SM::Offsets

SM-KeyAPI/include/utils.h

@@ -0,0 +1,89 @@
+#pragma once
+
+#define WIN32_LEAN_AND_MEAN
+#include <Windows.h>
+
+#include <string>
+#include <unordered_map>
+#include <algorithm>
+
+enum KeyState {
+	IDLE = 0,
+	JUST_PRESSED = 1,
+	HELD = 2,
+	RELEASED = 3
+};
+
+int GetVirtualKeyCode(const std::string& keyName) {
+    static const std::unordered_map<std::string, int> keyMap = {
+        {"0", '0'}, {"1", '1'}, {"2", '2'}, {"3", '3'}, {"4", '4'},
+        {"5", '5'}, {"6", '6'}, {"7", '7'}, {"8", '8'}, {"9", '9'},
+        {"enter", VK_RETURN}, {"space", VK_SPACE}, {"shift", VK_SHIFT},
+        {"ctrl", VK_CONTROL}, {"alt", VK_MENU}, {"tab", VK_TAB},
+        {"escape", VK_ESCAPE}, {"backspace", VK_BACK}, {"delete", VK_DELETE},
+        {"left", VK_LEFT}, {"right", VK_RIGHT}, {"Up", VK_UP}, {"Down", VK_DOWN},
+		{"f1", VK_F1}, {"f2", VK_F2}, {"f3", VK_F3}, {"f4", VK_F4}, {"f5", VK_F5},
+		{"f6", VK_F6}, {"f7", VK_F7}, {"f8", VK_F8}, {"f9", VK_F9}, {"f10", VK_F10},
+		{"f11", VK_F11}, {"f12", VK_F12},
+		{"num0", VK_NUMPAD0}, {"num1", VK_NUMPAD1}, {"num2", VK_NUMPAD2}, {"num3", VK_NUMPAD3}, {"num4", VK_NUMPAD4},
+		{"num5", VK_NUMPAD5}, {"num6", VK_NUMPAD6}, {"num7", VK_NUMPAD7}, {"num8", VK_NUMPAD8}, {"num9", VK_NUMPAD9},
+		{"num*", VK_MULTIPLY}, {"num+", VK_ADD}, {"num-", VK_SUBTRACT}, {"num.", VK_DECIMAL}, {"num/", VK_DIVIDE},
+		{"numLock", VK_NUMLOCK}, {"scrolllock", VK_SCROLL},
+		{"capsLock", VK_CAPITAL}, {"printscreen", VK_SNAPSHOT}, {"pause", VK_PAUSE},
+		{"insert", VK_INSERT}, {"home", VK_HOME}, {"end", VK_END}, {"pageup", VK_PRIOR}, {"pagedown", VK_NEXT},
+		{"lshift", VK_LSHIFT}, {"rshift", VK_RSHIFT}, {"lctrl", VK_LCONTROL}, {"rctrl", VK_RCONTROL},
+		{"lalt", VK_LMENU}, {"ralt", VK_RMENU}, {"lwin", VK_LWIN}, {"rwin", VK_RWIN}
+    };
+
+	std::string localKeyName = keyName;
+	std::transform(localKeyName.begin(), localKeyName.end(), localKeyName.begin(), ::tolower);
+
+	// If it's 1 letter, it's probably a character, just return the ASCII value
+	if (localKeyName.size() == 1) {
+		std::transform(localKeyName.begin(), localKeyName.end(), localKeyName.begin(), ::toupper);
+		return localKeyName[0];
+	}
+
+	auto it = keyMap.find(localKeyName);
+    if (it != keyMap.end()) {
+        return it->second;
+    }
+
+    return 0; // Invalid key
+}
+
+// perf: need to optimize this (use game InputManager?)
+KeyState getKeyState(const std::string& key) {
+	static std::unordered_map<int, bool> keyStates = {};
+
+	int keyCode = GetVirtualKeyCode(key);
+	SHORT keyState = GetAsyncKeyState(keyCode);
+	bool currentlyPressed = (keyState & 0x8000) != 0; // MSB is set if key is down
+
+	KeyState state = IDLE;
+	if (keyStates.find(keyCode) == keyStates.end()) {
+		if (currentlyPressed) {
+			keyStates[keyCode] = true;
+			state = JUST_PRESSED;
+		}
+	}
+	else {
+		if (currentlyPressed) {
+			if (keyStates[keyCode]) {
+				state = HELD;
+			}
+			else {
+				keyStates[keyCode] = true;
+				state = JUST_PRESSED;
+			}
+		}
+		else {
+			if (keyStates[keyCode]) {
+				keyStates[keyCode] = false;
+				state = RELEASED;
+			}
+		}
+	}
+
+	return state;
+}

SM-KeyAPI/src/dllmain.cpp

@@ -1,25 +1,70 @@
+#define WIN32_LEAN_AND_MEAN
+
 #include <Windows.h>
-#include <iostream>
+#include <optional>
 
-#include "log.h"
+#include <polyhook2/Detour/NatDetour.hpp>
+#include <polyhook2/IHook.hpp>
 
-// DllMain function which gets called when the DLL is loaded
-BOOL APIENTRY DllMain(HMODULE hModule, DWORD dwReason, LPVOID lpReserved) {
-	// RAII object to handle console creation/destruction
-	// Simply having this line here will allow you to print to the games console
-	// Don't worry about it, but if you're interested feel free to look in `utils.h`
-	// for the implementation
-	Log_t logRAII;
+#include <sm/offsets.h>
+#include <sm/lua.h>
+
+#include "utils.h"
+
+static int keyapi_getkeystate(lua_State* L) {
+	// Check that ScrapMechanic is the active window, we don't want any keyloggers made using SM...
+	if (GetForegroundWindow() != FindWindowA("CONTRAPTION_WINDOWS_CLASS", nullptr)) {
+		lua_pushinteger(L, 0);
+		return 1;
+	}
+
+	const char* key = luaL_checklstring(L, 1, nullptr);
+	lua_pushinteger(L, (int)getKeyState(std::string(key)));
+
+	return 1;
+}
+
+static std::optional<PLH::NatDetour> detour = std::nullopt;
+static uint64_t hkLoadLuaEnvTramp = NULL;
+// __int64 __fastcall loadLuaEnv(__int64 *luaVM, _QWORD **loadFuncs, int envFlag)
+NOINLINE uint64_t __cdecl hkLoadLuaEnv(uint64_t* luaVM, uint64_t** loadFuncs, int envFlag) {
+	auto oRes = PLH::FnCast(hkLoadLuaEnvTramp, &hkLoadLuaEnv)(luaVM, loadFuncs, envFlag);
+	lua_State* L = reinterpret_cast<lua_State*>(*luaVM);
+
+	if (!hasAllLuaFunctions())
+		return oRes;
 
+	if (!oRes && L) {
+		const int loadRes = luaL_loadstring(L, "unsafe_env.sm.keyapi_injected = true");
+		if (!loadRes)
+			lua_pcall(L, 0, 0, 0);
+
+		lua_getglobal(L, "unsafe_env");
+		lua_getfield(L, -1, "sm");
+		lua_pushcclosure(L, (lua_CFunction*)keyapi_getkeystate, 0);
+		lua_setfield(L, -2, "getKeyState");
+
+		return loadRes;
+	}
+
+	return oRes;
+}
+
+BOOL APIENTRY DllMain(HMODULE hModule, DWORD dwReason, LPVOID lpReserved) {
 	if (dwReason == DLL_PROCESS_ATTACH) {
 		DisableThreadLibraryCalls(hModule);
 		std::cout << "DLL_PROCESS_ATTACH" << std::endl;
-		// Setup, hook, etc here
+
+		detour.emplace((uint64_t)SM::Offsets::Rebased::loadLuaEnv, (uint64_t)&hkLoadLuaEnv, &hkLoadLuaEnvTramp);
+		detour.value().hook();
+		std::cout << "loadLuaEnv " << std::hex << (uint64_t)SM::Offsets::Rebased::loadLuaEnv << " hooked!" << std::endl;
 	}
 
 	if (dwReason == DLL_PROCESS_DETACH) {
 		std::cout << "DLL_PROCESS_DETACH" << std::endl;
-		// Cleanup, unhooking, etc here
+
+		detour.value().unHook();
+		std::cout << "loadLuaEnv " << std::hex << (uint64_t)SM::Offsets::Rebased::loadLuaEnv << " unhooked!" << std::endl;
 	}
 
 	return TRUE;

SM-KeyAPI/vcpkg-configuration.json

@@ -0,0 +1,14 @@
+{
+  "default-registry": {
+    "kind": "git",
+    "baseline": "d033613d9021107e4a7b52c5fac1f87ae8a6fcc6",
+    "repository": "https://github.com/microsoft/vcpkg"
+  },
+  "registries": [
+    {
+      "kind": "artifact",
+      "location": "https://github.com/microsoft/vcpkg-ce-catalog/archive/refs/heads/main.zip",
+      "name": "microsoft"
+    }
+  ]
+}

SM-KeyAPI/vcpkg.json

@@ -0,0 +1,5 @@
+{
+  "dependencies": [
+    "polyhook2"
+  ]
+}