1  #pragma once
  2  
  3  #define WIN32_LEAN_AND_MEAN
  4  #include <Windows.h>
  5  #include <shellapi.h>
  6  #include <sddl.h>
  7  #include <shldisp.h>
  8  #include <shlobj.h>
  9  #include <exdisp.h>
 10  #include <atlbase.h>
 11  #include <stdlib.h>
 12  #include <comdef.h>
 13  
 14  #include <winrt/base.h>
 15  #include <winrt/Windows.Foundation.Collections.h>
 16  
 17  #include <string>
 18  #include <filesystem>
 19  
 20  #include <common/logger/logger.h>
 21  #include <common/utils/winapi_error.h>
 22  #include <common/utils/process_path.h>
 23  #include <common/utils/processApi.h>
 24  
 25  namespace
 26  {
 27      inline std::wstring GetErrorString(HRESULT handle)
 28      {
 29          _com_error err(handle);
 30          return err.ErrorMessage();
 31      }
 32  
 33      inline bool FindDesktopFolderView(REFIID riid, void** ppv)
 34      {
 35          CComPtr<IShellWindows> spShellWindows;
 36          auto result = spShellWindows.CoCreateInstance(CLSID_ShellWindows);
 37          if (result != S_OK || spShellWindows == nullptr)
 38          {
 39              Logger::warn(L"Failed to create instance. {}", GetErrorString(result));
 40              return false;
 41          }
 42  
 43          CComVariant vtLoc(CSIDL_DESKTOP);
 44          CComVariant vtEmpty;
 45          long lhwnd;
 46          CComPtr<IDispatch> spdisp;
 47          result = spShellWindows->FindWindowSW(
 48              &vtLoc, &vtEmpty, SWC_DESKTOP, &lhwnd, SWFO_NEEDDISPATCH, &spdisp);
 49  
 50          if (result != S_OK || spdisp == nullptr)
 51          {
 52              Logger::warn(L"Failed to find the window. {}", GetErrorString(result));
 53              return false;
 54          }
 55  
 56          CComPtr<IShellBrowser> spBrowser;
 57          result = CComQIPtr<IServiceProvider>(spdisp)->QueryService(SID_STopLevelBrowser,
 58                                                                     IID_PPV_ARGS(&spBrowser));
 59          if (result != S_OK || spBrowser == nullptr)
 60          {
 61              Logger::warn(L"Failed to query service. {}", GetErrorString(result));
 62              return false;
 63          }
 64  
 65          CComPtr<IShellView> spView;
 66          result = spBrowser->QueryActiveShellView(&spView);
 67          if (result != S_OK || spView == nullptr)
 68          {
 69              Logger::warn(L"Failed to query active shell window. {}", GetErrorString(result));
 70              return false;
 71          }
 72  
 73          result = spView->QueryInterface(riid, ppv);
 74          if (result != S_OK || ppv == nullptr || *ppv == nullptr)
 75          {
 76              Logger::warn(L"Failed to query interface. {}", GetErrorString(result));
 77              return false;
 78          }
 79  
 80          return true;
 81      }
 82  
 83      inline bool GetDesktopAutomationObject(REFIID riid, void** ppv)
 84      {
 85          CComPtr<IShellView> spsv;
 86  
 87          // Desktop may not be available on startup
 88          auto attempts = 5;
 89          for (auto i = 1; i <= attempts; i++)
 90          {
 91              if (FindDesktopFolderView(IID_PPV_ARGS(&spsv)))
 92              {
 93                  break;
 94              }
 95  
 96              Logger::warn(L"FindDesktopFolderView() failed attempt {}", i);
 97  
 98              if (i == attempts)
 99              {
100                  Logger::warn(L"FindDesktopFolderView() max attempts reached");
101                  return false;
102              }
103  
104              Sleep(3000);
105          }
106  
107          CComPtr<IDispatch> spdispView;
108          auto result = spsv->GetItemObject(SVGIO_BACKGROUND, IID_PPV_ARGS(&spdispView));
109          if (result != S_OK)
110          {
111              Logger::warn(L"GetItemObject() failed. {}", GetErrorString(result));
112              return false;
113          }
114  
115          result = spdispView->QueryInterface(riid, ppv);
116          if (result != S_OK)
117          {
118              Logger::warn(L"QueryInterface() failed. {}", GetErrorString(result));
119              return false;
120          }
121  
122          return true;
123      }
124  
125      inline bool ShellExecuteFromExplorer(
126          PCWSTR pszFile,
127          PCWSTR pszParameters = nullptr,
128          PCWSTR workingDir = L"")
129      {
130          CComPtr<IShellFolderViewDual> spFolderView;
131          if (!GetDesktopAutomationObject(IID_PPV_ARGS(&spFolderView)))
132          {
133              return false;
134          }
135  
136          CComPtr<IDispatch> spdispShell;
137          auto result = spFolderView->get_Application(&spdispShell);
138          if (result != S_OK)
139          {
140              Logger::warn(L"get_Application() failed. {}", GetErrorString(result));
141              return false;
142          }
143  
144          CComQIPtr<IShellDispatch2>(spdispShell)
145              ->ShellExecuteW(CComBSTR(pszFile),
146                              CComVariant(pszParameters ? pszParameters : L""),
147                              CComVariant(workingDir),
148                              CComVariant(L""),
149                              CComVariant(SW_SHOWNORMAL));
150  
151          return true;
152      }
153  }
154  
155  // Returns true if the current process is running with elevated privileges
156  inline bool is_process_elevated(const bool use_cached_value = true)
157  {
158      auto detection_func = []() {
159          HANDLE token = nullptr;
160          bool elevated = false;
161  
162          if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &token))
163          {
164              TOKEN_ELEVATION elevation;
165              DWORD size;
166              if (GetTokenInformation(token, TokenElevation, &elevation, sizeof(elevation), &size))
167              {
168                  elevated = (elevation.TokenIsElevated != 0);
169              }
170          }
171  
172          if (token)
173          {
174              CloseHandle(token);
175          }
176  
177          return elevated;
178      };
179      static const bool cached_value = detection_func();
180      return use_cached_value ? cached_value : detection_func();
181  }
182  
183  // Drops the elevated privileges if present
184  inline bool drop_elevated_privileges()
185  {
186      HANDLE token = nullptr;
187      if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_DEFAULT | WRITE_OWNER, &token))
188      {
189          return false;
190      }
191  
192      PSID medium_sid = NULL;
193      if (!::ConvertStringSidToSid(SDDL_ML_MEDIUM, &medium_sid))
194      {
195          return false;
196      }
197  
198      TOKEN_MANDATORY_LABEL label = { 0 };
199      label.Label.Attributes = SE_GROUP_INTEGRITY;
200      label.Label.Sid = medium_sid;
201      DWORD size = static_cast<DWORD>(sizeof(TOKEN_MANDATORY_LABEL) + ::GetLengthSid(medium_sid));
202  
203      BOOL result = SetTokenInformation(token, TokenIntegrityLevel, &label, size);
204      LocalFree(medium_sid);
205      CloseHandle(token);
206  
207      return result;
208  }
209  
210  // Run command as different user, returns true if succeeded
211  inline HANDLE run_as_different_user(const std::wstring& file, const std::wstring& params, const wchar_t* workingDir = nullptr, const bool showWindow = true)
212  {
213      Logger::info(L"run_elevated with params={}", params);
214      SHELLEXECUTEINFOW exec_info = { 0 };
215      exec_info.cbSize = sizeof(SHELLEXECUTEINFOW);
216      exec_info.lpVerb = L"runAsUser";
217      exec_info.lpFile = file.c_str();
218      exec_info.lpParameters = params.c_str();
219      exec_info.hwnd = 0;
220      exec_info.fMask = SEE_MASK_NOCLOSEPROCESS;
221      exec_info.lpDirectory = workingDir;
222      exec_info.hInstApp = 0;
223      if (showWindow)
224      {
225          exec_info.nShow = SW_SHOWDEFAULT;
226      }
227      else
228      {
229          // might have limited success, but only option with ShellExecuteExW
230          exec_info.nShow = SW_HIDE;
231      }
232  
233      return ShellExecuteExW(&exec_info) ? exec_info.hProcess : nullptr;
234  }
235  
236  // Run command as elevated user, returns true if succeeded
237  inline HANDLE run_elevated(const std::wstring& file, const std::wstring& params, const wchar_t* workingDir = nullptr, const bool showWindow = true)
238  {
239      Logger::info(L"run_elevated with params={}", params);
240      SHELLEXECUTEINFOW exec_info = { 0 };
241      exec_info.cbSize = sizeof(SHELLEXECUTEINFOW);
242      exec_info.lpVerb = L"runas";
243      exec_info.lpFile = file.c_str();
244      exec_info.lpParameters = params.c_str();
245      exec_info.hwnd = 0;
246      exec_info.fMask = SEE_MASK_NOCLOSEPROCESS;
247      exec_info.lpDirectory = workingDir;
248      exec_info.hInstApp = 0;
249  
250      if (showWindow)
251      {
252          exec_info.nShow = SW_SHOWDEFAULT;
253      }
254      else
255      {
256          // might have limited success, but only option with ShellExecuteExW
257          exec_info.nShow = SW_HIDE;
258      }
259  
260      BOOL result = ShellExecuteExW(&exec_info);
261  
262      return result  ? exec_info.hProcess : nullptr;
263  }
264  
265  // Run command as non-elevated user, returns true if succeeded, puts the process id into returnPid if returnPid != NULL
266  inline bool run_non_elevated(const std::wstring& file, const std::wstring& params, DWORD* returnPid, const wchar_t* workingDir = nullptr, const bool showWindow = true)
267  {
268      Logger::info(L"run_non_elevated with params={}", params);
269      auto executable_args = L"\"" + file + L"\"";
270      if (!params.empty())
271      {
272          executable_args += L" " + params;
273      }
274  
275      HWND hwnd = GetShellWindow();
276      if (!hwnd)
277      {
278          if (GetLastError() == ERROR_SUCCESS)
279          {
280              Logger::warn(L"GetShellWindow() returned null. Shell window is not available");
281          }
282          else
283          {
284              Logger::error(L"GetShellWindow() failed. {}", get_last_error_or_default(GetLastError()));
285          }
286  
287          return false;
288      }
289      DWORD pid;
290      GetWindowThreadProcessId(hwnd, &pid);
291  
292      winrt::handle process{ OpenProcess(PROCESS_CREATE_PROCESS, FALSE, pid) };
293      if (!process)
294      {
295          Logger::error(L"OpenProcess() failed. {}", get_last_error_or_default(GetLastError()));
296          return false;
297      }
298  
299      SIZE_T size = 0;
300  
301      InitializeProcThreadAttributeList(nullptr, 1, 0, &size);
302      auto pproc_buffer = std::make_unique<char[]>(size);
303      auto pptal = reinterpret_cast<PPROC_THREAD_ATTRIBUTE_LIST>(pproc_buffer.get());
304      if (!pptal)
305      {
306          Logger::error(L"pptal failed to initialize. {}", get_last_error_or_default(GetLastError()));
307          return false;
308      }
309  
310      if (!InitializeProcThreadAttributeList(pptal, 1, 0, &size))
311      {
312          Logger::error(L"InitializeProcThreadAttributeList() failed. {}", get_last_error_or_default(GetLastError()));
313          return false;
314      }
315  
316      HANDLE process_handle = process.get();
317      if (!UpdateProcThreadAttribute(pptal,
318                                     0,
319                                     PROC_THREAD_ATTRIBUTE_PARENT_PROCESS,
320                                     &process_handle,
321                                     sizeof(process_handle),
322                                     nullptr,
323                                     nullptr))
324      {
325          Logger::error(L"UpdateProcThreadAttribute() failed. {}", get_last_error_or_default(GetLastError()));
326          return false;
327      }
328  
329      STARTUPINFOEX siex = { 0 };
330      siex.lpAttributeList = pptal;
331      siex.StartupInfo.cb = sizeof(siex);
332      PROCESS_INFORMATION pi = { 0 };
333      auto dwCreationFlags = EXTENDED_STARTUPINFO_PRESENT;
334  
335      if (!showWindow)
336      {
337          siex.StartupInfo.dwFlags = STARTF_USESHOWWINDOW;
338          siex.StartupInfo.wShowWindow = SW_HIDE;
339          dwCreationFlags = CREATE_NO_WINDOW;
340      }
341  
342      auto succeeded = CreateProcessW(file.c_str(),
343                                      &executable_args[0],
344                                      nullptr,
345                                      nullptr,
346                                      FALSE,
347                                      dwCreationFlags,
348                                      nullptr,
349                                      workingDir,
350                                      &siex.StartupInfo,
351                                      &pi);
352      if (succeeded)
353      {
354          if (pi.hProcess)
355          {
356              if (returnPid)
357              {
358                  *returnPid = GetProcessId(pi.hProcess);
359              }
360  
361              CloseHandle(pi.hProcess);
362          }
363          if (pi.hThread)
364          {
365              CloseHandle(pi.hThread);
366          }
367      }
368      else
369      {
370          Logger::error(L"CreateProcessW() failed. {}", get_last_error_or_default(GetLastError()));
371      }
372  
373      return succeeded;
374  }
375  
376  inline bool RunNonElevatedEx(const std::wstring& file, const std::wstring& params, const std::wstring& working_dir)
377  {
378      bool success = false;
379      HRESULT co_init = E_FAIL;
380      try
381      {
382          co_init = CoInitialize(nullptr);
383          success = ShellExecuteFromExplorer(file.c_str(), params.c_str(), working_dir.c_str());
384      }
385      catch (...)
386      {
387      }
388      if (SUCCEEDED(co_init))
389      {
390          CoUninitialize();
391      }
392  
393      return success;
394  }
395  
396  struct ProcessInfo
397  {
398      wil::unique_process_handle processHandle;
399      DWORD processID = {};
400  };
401  
402  inline std::optional<ProcessInfo> RunNonElevatedFailsafe(const std::wstring& file, const std::wstring& params, const std::wstring& working_dir, DWORD handleAccess = 0)
403  {
404      bool launched = RunNonElevatedEx(file, params, working_dir);
405      if (!launched)
406      {
407          Logger::warn(L"RunNonElevatedEx() failed. Trying fallback");
408          std::wstring action_runner_path = get_module_folderpath() + L"\\PowerToys.ActionRunner.exe";
409          std::wstring newParams = fmt::format(L"-run-non-elevated -target \"{}\" {}", file, params);
410          launched = run_non_elevated(action_runner_path, newParams, nullptr, working_dir.c_str());
411          if (launched)
412          {
413              Logger::trace(L"Started {}", file);
414          }
415          else
416          {
417              Logger::warn(L"Failed to start {}", file);
418              return std::nullopt;
419          }
420      }
421  
422      auto handles = getProcessHandlesByName(std::filesystem::path{ file }.filename().wstring(), PROCESS_QUERY_INFORMATION | SYNCHRONIZE | handleAccess);
423  
424      if (handles.empty())
425          return std::nullopt;
426  
427      ProcessInfo result;
428      result.processID = GetProcessId(handles[0].get());
429      result.processHandle = std::move(handles[0]);
430  
431      return result;
432  }
433  
434  // Run command with the same elevation, returns true if succeeded
435  inline bool run_same_elevation(const std::wstring& file, const std::wstring& params, DWORD* returnPid, const wchar_t* workingDir = nullptr)
436  {
437      auto executable_args = L"\"" + file + L"\"";
438      if (!params.empty())
439      {
440          executable_args += L" " + params;
441      }
442  
443      STARTUPINFO si = { sizeof(STARTUPINFO) };
444      PROCESS_INFORMATION pi = { 0 };
445  
446      auto succeeded = CreateProcessW(file.c_str(),
447                                      &executable_args[0],
448                                      nullptr,
449                                      nullptr,
450                                      FALSE,
451                                      0,
452                                      nullptr,
453                                      workingDir,
454                                      &si,
455                                      &pi);
456  
457      if (succeeded)
458      {
459          if (pi.hProcess)
460          {
461              if (returnPid)
462              {
463                  *returnPid = GetProcessId(pi.hProcess);
464              }
465  
466              CloseHandle(pi.hProcess);
467          }
468  
469          if (pi.hThread)
470          {
471              CloseHandle(pi.hThread);
472          }
473      }
474      return succeeded;
475  }
476  
477  // Returns true if the current process is running from administrator account
478  // The function returns true in case of error since we want to return false
479  // only in case of a positive verification that the user is not an admin.
480  inline bool check_user_is_admin()
481  {
482      auto freeMemory = [](PSID pSID, PTOKEN_GROUPS pGroupInfo) {
483          if (pSID)
484          {
485              FreeSid(pSID);
486          }
487          if (pGroupInfo)
488          {
489              GlobalFree(pGroupInfo);
490          }
491      };
492  
493      HANDLE hToken;
494      DWORD dwSize = 0;
495      PTOKEN_GROUPS pGroupInfo;
496      SID_IDENTIFIER_AUTHORITY SIDAuth = SECURITY_NT_AUTHORITY;
497      PSID pSID = NULL;
498  
499      // Open a handle to the access token for the calling process.
500      if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken))
501      {
502          return true;
503      }
504  
505      // Call GetTokenInformation to get the buffer size.
506      if (!GetTokenInformation(hToken, TokenGroups, NULL, dwSize, &dwSize))
507      {
508          if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
509          {
510              return true;
511          }
512      }
513  
514      // Allocate the buffer.
515      pGroupInfo = static_cast<PTOKEN_GROUPS>(GlobalAlloc(GPTR, dwSize));
516  
517      // Call GetTokenInformation again to get the group information.
518      if (!GetTokenInformation(hToken, TokenGroups, pGroupInfo, dwSize, &dwSize))
519      {
520          freeMemory(pSID, pGroupInfo);
521          return true;
522      }
523  
524      // Create a SID for the BUILTIN\Administrators group.
525      if (!AllocateAndInitializeSid(&SIDAuth, 2, SECURITY_BUILTIN_DOMAIN_RID, DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &pSID))
526      {
527          freeMemory(pSID, pGroupInfo);
528          return true;
529      }
530  
531      // Loop through the group SIDs looking for the administrator SID.
532      for (DWORD i = 0; i < pGroupInfo->GroupCount; ++i)
533      {
534          if (EqualSid(pSID, pGroupInfo->Groups[i].Sid))
535          {
536              freeMemory(pSID, pGroupInfo);
537              return true;
538          }
539      }
540  
541      freeMemory(pSID, pGroupInfo);
542      return false;
543  }
544  
545  inline bool IsProcessOfWindowElevated(HWND window)
546  {
547      DWORD pid = 0;
548      GetWindowThreadProcessId(window, &pid);
549      if (!pid)
550      {
551          return false;
552      }
553  
554      wil::unique_handle hProcess{ OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION,
555                                               FALSE,
556                                               pid) };
557  
558      wil::unique_handle token;
559  
560      if (OpenProcessToken(hProcess.get(), TOKEN_QUERY, &token))
561      {
562          TOKEN_ELEVATION elevation;
563          DWORD size;
564          if (GetTokenInformation(token.get(), TokenElevation, &elevation, sizeof(elevation), &size))
565          {
566              return elevation.TokenIsElevated != 0;
567          }
568      }
569      return false;
570  }