DLL injection path bug fix

Author: leonav-unizar

include/keyreaper/config.h

@@ -9,7 +9,7 @@ class Config {
  public:
   static Config& Instance();  // Singleton accessor
 
-  error_handling::ProgramResult Load(const std::string& filename);
+  error_handling::ProgramResult Load(const std::wstring& filename);
 
   std::wstring GetKeyExtractorDLLPath() const;
 
@@ -20,7 +20,7 @@ class Config {
   Config(const Config&) = delete;
   Config& operator=(const Config&) = delete;
 
-  std::string key_extractor_dll_;
+  std::wstring key_extractor_dll_;
 };
 
 #endif  // CONFIG_H_

src/config.cc

@@ -11,52 +11,50 @@ Config& Config::Instance() {
   return instance;
 }
 
-wstring StringToWString(const string& str) {
-  if (str.empty()) return wstring();
-
-  int size_needed = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, nullptr, 0);
-  if (size_needed == 0) return wstring();  // Conversion failed
-
-  wstring wstr(size_needed, 0);
-  MultiByteToWideChar(CP_UTF8, 0, str.c_str(), -1, &wstr[0], size_needed);
-
-  return wstr;
+wstring ExePath() {
+  wchar_t buffer[MAX_PATH] = { 0 };
+  GetModuleFileNameW(NULL, buffer, MAX_PATH);
+  wstring fullPath(buffer);
+  wstring::size_type pos = fullPath.find_last_of(L"\\/");
+  return fullPath.substr(0, pos);
 }
 
-
-bool FileExists(string filename) {
+bool FileExists(wstring filename) {
   ifstream file_check(filename);
   if (file_check) return true;
   else return false;
 }
 
-ProgramResult Config::Load(const string& filename) {
+ProgramResult Config::Load(const wstring& filename) {
   if (!FileExists(filename)) return ErrorResult("Could not find config file");
 
   toml::table config;
   try {
     config = toml::parse_file(filename);
 
   } catch (const toml::parse_error& err) {
-    cerr << err << endl;
-    return ErrorResult("Failed to parse config file " + filename);
+    wcerr << L"Error while parsing " << filename; 
+    cerr << ": " << err << endl;
+    return ErrorResult("Failed to parse config file ");
   }
 
+  auto current_path = ExePath();
+
 #if _WIN64
-  string default_dll = "injectable_server_x64.dll";
+  wstring default_dll = current_path + L"\\injectable_server_x64.dll";
 #else
-  string default_dll = "injectable_server_x86.dll";
+  wstring default_dll = current_path + L"\\injectable_server_x86.dll";
 #endif
   key_extractor_dll_ = config["Settings"]["key_extractor_dll_path"].value_or(default_dll);
   if (!FileExists(key_extractor_dll_)) {
-    cerr << "Could not find a default\n";
+    cerr << "Default DLL not set in the config file\n";
     if (!FileExists(default_dll)) return ErrorResult("Could not find a valid DLL");
-    cerr << "Specified path to DLL does not exist. Defaulting to " << default_dll << endl;
+    wcerr << "Specified path to DLL does not exist. Defaulting to " << default_dll << endl;
   }
 
   return OkResult("Configuration successfully loaded");
 }
 
 std::wstring Config::GetKeyExtractorDLLPath() const {
-  return StringToWString(key_extractor_dll_);
+  return key_extractor_dll_;
 }

src/injection/injector.cc

@@ -14,6 +14,14 @@ using namespace error_handling;
 
 namespace injection {
 
+wstring DllName(wstring dll_path) {
+  std::wstring::size_type pos = dll_path.find_last_of(L"\\/");
+  if (pos == std::wstring::npos)
+    return dll_path; // No path separator found, return entire string
+
+  return dll_path.substr(pos + 1);
+}
+
 bool IsDLLLoadedOnProcess(DWORD pid, wstring w_dll_path, HANDLE process_handle) {
   if (!process_handle) return false;
 
@@ -23,7 +31,7 @@ bool IsDLLLoadedOnProcess(DWORD pid, wstring w_dll_path, HANDLE process_handle)
     for (size_t i = 0; i < (bytes_needed / sizeof(HMODULE)); i++) {
       WCHAR module_name[MAX_PATH];
       if (GetModuleBaseNameW(process_handle, module_array[i], module_name, sizeof(module_name) / sizeof(WCHAR))) {
-        if (_wcsicmp(module_name, w_dll_path.c_str()) == 0) {
+        if (_wcsicmp(module_name, DllName(w_dll_path).c_str()) == 0) {
           printf(" [i] DLL found in process\n");
           return true;
         }

src/main.cc

@@ -56,9 +56,9 @@ std::string GetChoices(const std::map<std::string, EnumClass>& options_map, cons
 }
 
 bool LoadConfig() {
-  string config_file = "config.toml";
+  wstring config_file = L"config.toml";
   if (!fs::exists(config_file)) {
-    cout << "[!] Config file not found, creating a new one with the name: " << config_file << endl;
+    wcout << "[!] Config file not found, creating a new one with the name: " << config_file << endl;
     ofstream empty_file(config_file);
   }