1  #include "pch.h"
  2  #include "NativeMethods.h"
  3  #include "FileLocksmith.h"
  4  #include "../FileLocksmithLib/Constants.h"
  5  
  6  namespace winrt::PowerToys::FileLocksmithLib::Interop::implementation
  7  {
  8  
  9  #pragma region HelperMethods
 10      std::wstring executable_path()
 11      {
 12          return pid_to_full_path(GetCurrentProcessId());
 13      }
 14      std::wstring paths_file()
 15      {
 16          std::wstring path{ PowerToys::Interop::Constants::AppDataPath() };
 17          path += L"\\";
 18          path += constants::nonlocalizable::PowerToyName;
 19          path += L"\\";
 20          path += constants::nonlocalizable::LastRunPath;
 21          return path;
 22      }
 23  
 24  #pragma endregion
 25  
 26      com_array<winrt::PowerToys::FileLocksmithLib::Interop::ProcessResult> NativeMethods::FindProcessesRecursive(array_view<hstring const> paths)
 27      {
 28          std::vector<std::wstring> paths_cpp{ paths.begin(), paths.end() };
 29  
 30          auto result_cpp = find_processes_recursive(paths_cpp);
 31          const auto result_size = static_cast<int>(result_cpp.size());
 32  
 33          std::vector<ProcessResult> result;
 34  
 35          if (result_size == 0)
 36          {
 37              return com_array<ProcessResult>();
 38          }
 39  
 40          for (int i = 0; i < result_size; i++)
 41          {
 42              result.push_back(ProcessResult
 43              {
 44                  hstring { result_cpp[i].name },
 45                  result_cpp[i].pid,
 46                  hstring{ result_cpp[i].user },
 47                      winrt::com_array<hstring>
 48                  {
 49                      result_cpp[i].files.begin(), result_cpp[i].files.end()
 50                  }
 51              });
 52          }
 53  
 54          return com_array<ProcessResult>{ result.begin(), result.end() };
 55      }
 56  
 57      hstring NativeMethods::PidToFullPath(uint32_t pid)
 58      {
 59          return hstring{ pid_to_full_path(pid) };
 60      }
 61  
 62      com_array<hstring> NativeMethods::ReadPathsFromFile()
 63      {
 64          std::ifstream stream(paths_file());
 65  
 66          std::vector<std::wstring> result_cpp;
 67          std::wstring line;
 68  
 69          bool finished = false;
 70  
 71          while (!finished)
 72          {
 73              WCHAR ch{};
 74              // We have to read data like this
 75              if (!stream.read(reinterpret_cast<char*>(&ch), 2))
 76              {
 77                  finished = true;
 78              }
 79              else if (ch == L'\n')
 80              {
 81                  if (line.empty())
 82                  {
 83                      finished = true;
 84                  }
 85                  else
 86                  {
 87                      result_cpp.push_back(line);
 88                      line = {};
 89                  }
 90              }
 91              else
 92              {
 93                  line += ch;
 94              }
 95          }
 96          return com_array<hstring>{ result_cpp.begin(), result_cpp.end() };
 97      }
 98   
 99      bool NativeMethods::StartAsElevated(array_view<hstring const> paths)
100      {
101          std::ofstream stream(paths_file());
102          const WCHAR newline = L'\n';
103  
104          for (uint32_t i = 0; i < paths.size(); i++)
105          {
106              std::wstring path_cpp{ paths[i] };
107              stream.write(reinterpret_cast<const char*>(path_cpp.c_str()), path_cpp.size() * sizeof(WCHAR));
108              stream.write(reinterpret_cast<const char*>(&newline), sizeof(WCHAR));
109          }
110  
111          stream.write(reinterpret_cast<const char*>(&newline), sizeof(WCHAR));
112  
113          if (!stream)
114          {
115              return false;
116          }
117  
118          stream.close();
119  
120          auto exec_path = executable_path();
121  
122          SHELLEXECUTEINFOW exec_info{};
123          exec_info.cbSize = sizeof(exec_info);
124          exec_info.fMask = SEE_MASK_NOCLOSEPROCESS;
125          exec_info.hwnd = NULL;
126          exec_info.lpVerb = L"runas";
127          exec_info.lpFile = exec_path.c_str();
128          exec_info.lpParameters = L"--elevated";
129          exec_info.lpDirectory = NULL;
130          exec_info.nShow = SW_SHOW;
131          exec_info.hInstApp = NULL;
132  
133          if (ShellExecuteExW(&exec_info))
134          {
135              CloseHandle(exec_info.hProcess);
136              return true;
137          }
138  
139          return false;
140      }
141  
142      /* Adapted from "https://learn.microsoft.com/windows/win32/secauthz/enabling-and-disabling-privileges-in-c--" */
143      bool NativeMethods::SetDebugPrivilege()
144      {
145          HANDLE hToken;
146          TOKEN_PRIVILEGES tp{};
147          LUID luid;
148  
149          if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken) != 0)
150          {
151              if (!LookupPrivilegeValue(
152                      NULL, // lookup privilege on local system
153                      SE_DEBUG_NAME, // privilege to lookup
154                      &luid)) // receives LUID of privilege
155              {
156                  CloseHandle(hToken);
157                  return false;
158              }
159              tp.PrivilegeCount = 1;
160              tp.Privileges[0].Luid = luid;
161              tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
162  
163              if (!AdjustTokenPrivileges(
164                      hToken,
165                      FALSE,
166                      &tp,
167                      sizeof(TOKEN_PRIVILEGES),
168                      NULL,
169                      NULL))
170              {
171                  CloseHandle(hToken);
172                  return false;
173              }
174  
175              if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
176              {
177                  CloseHandle(hToken);
178                  return false;
179              }
180  
181              CloseHandle(hToken);
182              return true;
183          }
184          return false;
185      }
186   
187      // adapted from common/utils/elevation.h. No need to bring all dependencies to this project, though.
188      // TODO: Make elevation.h lighter so that this function can be used without bringing dependencies like spdlog in.
189      bool NativeMethods::IsProcessElevated()
190      {
191          HANDLE token = nullptr;
192          bool elevated = false;
193          if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token))
194          {
195              TOKEN_ELEVATION elevation{};
196              DWORD size;
197              if (GetTokenInformation(token, TokenElevation, &elevation, sizeof(elevation), &size))
198              {
199                  elevated = (elevation.TokenIsElevated != 0);
200              }
201          }
202          if (token)
203          {
204              CloseHandle(token);
205          }
206          return elevated;
207      }
208  }