From 5515b65ca26cba48ee887cb8a0d28ad4b172819b Mon Sep 17 00:00:00 2001
From: WindyValley <2519087323@qq.com>
Date: Sun, 13 Oct 2024 11:25:15 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E9=BB=98=E8=AE=A4=E5=90=AF=E5=8A=A8?=
 =?UTF-8?q?=E4=B8=8D=E8=AF=B7=E6=B1=82=E7=AE=A1=E7=90=86=E5=91=98=E6=9D=83?=
 =?UTF-8?q?=E9=99=90=EF=BC=8C=E5=BF=85=E8=A6=81=E6=97=B6=E4=BD=BF=E7=94=A8?=
 =?UTF-8?q?=E7=AE=A1=E7=90=86=E5=91=98=E6=9D=83=E9=99=90=E9=87=8D=E5=90=AF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 写 HKCU 的命令不需要管理员权限，工程文件里的 RequireAdministrator
  改为 AsInvoker
- 调用 uninstall 和 install/CustomInstall 前判断一下进程的状态，
  如果不是管理就提权重启
---
 WeaselSetup/WeaselSetup.cpp     | 57 +++++++++++++++++++++++++++++++--
 WeaselSetup/WeaselSetup.vcxproj |  4 +--
 WeaselSetup/xmake.lua           |  2 +-
 3 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/WeaselSetup/WeaselSetup.cpp b/WeaselSetup/WeaselSetup.cpp
index 259af11f7..c330d2f77 100644
--- a/WeaselSetup/WeaselSetup.cpp
+++ b/WeaselSetup/WeaselSetup.cpp
@@ -14,6 +14,8 @@
 CAppModule _Module;
 
 static int Run(LPTSTR lpCmdLine);
+static bool IsProcAdmin();
+static int RestartAsAdmin(LPTSTR lpCmdLine);
 
 int WINAPI _tWinMain(HINSTANCE hInstance,
                      HINSTANCE /*hPrevInstance*/,
@@ -138,8 +140,12 @@ static int Run(LPTSTR lpCmdLine) {
   constexpr bool silent = true;
   constexpr bool old_ime_support = false;
   bool uninstalling = !wcscmp(L"/u", lpCmdLine);
-  if (uninstalling)
-    return uninstall(silent);
+  if (uninstalling) {
+    if (IsProcAdmin())
+      return uninstall(silent);
+    else
+      RestartAsAdmin(lpCmdLine);
+  }
 
   if (!wcscmp(L"/ls", lpCmdLine)) {
     return SetRegKeyValue(HKEY_CURRENT_USER, L"Software\\Rime\\weasel",
@@ -177,6 +183,12 @@ static int Run(LPTSTR lpCmdLine) {
     return SetRegKeyValue(HKEY_CURRENT_USER, L"Software\\Rime\\weasel",
                           L"UpdateChannel", L"release", REG_SZ);
   }
+
+  if (!IsProcAdmin())
+  {
+    return RestartAsAdmin(lpCmdLine);
+  }
+
   bool hans = !wcscmp(L"/s", lpCmdLine);
   if (hans)
     return install(false, silent, old_ime_support);
@@ -186,3 +198,44 @@ static int Run(LPTSTR lpCmdLine) {
   bool installing = !wcscmp(L"/i", lpCmdLine);
   return CustomInstall(installing);
 }
+
+// https://learn.microsoft.com/zh-cn/windows/win32/api/securitybaseapi/nf-securitybaseapi-checktokenmembership
+bool IsProcAdmin() {
+  BOOL b = FALSE;
+  SID_IDENTIFIER_AUTHORITY NtAuthority = SECURITY_NT_AUTHORITY;
+  PSID AdministratorsGroup;
+  b = AllocateAndInitializeSid(&NtAuthority, 2, SECURITY_BUILTIN_DOMAIN_RID,
+                               DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0,
+                               &AdministratorsGroup);
+
+  if (b) {
+    if (!CheckTokenMembership(NULL, AdministratorsGroup, &b)) {
+      b = FALSE;
+    }
+    FreeSid(AdministratorsGroup);
+  }
+
+  return (b);
+}
+
+int RestartAsAdmin(LPTSTR lpCmdLine) {
+  SHELLEXECUTEINFO execInfo{0};
+  TCHAR path[MAX_PATH];
+  GetModuleFileName(GetModuleHandle(NULL), path, _countof(path));
+  execInfo.lpFile = path;
+  execInfo.lpParameters = lpCmdLine;
+  execInfo.lpVerb = _T("runas");
+  execInfo.cbSize = sizeof(execInfo);
+  execInfo.nShow = SW_SHOWNORMAL;
+  execInfo.fMask = SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS;
+  execInfo.hwnd = NULL;
+  execInfo.hProcess = NULL;
+  if (::ShellExecuteEx(&execInfo) && execInfo.hProcess != NULL) {
+    ::WaitForSingleObject(execInfo.hProcess, INFINITE);
+    DWORD dwExitCode = 0;
+    ::GetExitCodeProcess(execInfo.hProcess, &dwExitCode);
+    ::CloseHandle(execInfo.hProcess);
+    return dwExitCode;
+  }
+  return -1;
+}
diff --git a/WeaselSetup/WeaselSetup.vcxproj b/WeaselSetup/WeaselSetup.vcxproj
index 5a2c76763..632bb0114 100644
--- a/WeaselSetup/WeaselSetup.vcxproj
+++ b/WeaselSetup/WeaselSetup.vcxproj
@@ -108,7 +108,7 @@
       <GenerateDebugInformation>true</GenerateDebugInformation>
       <AdditionalDependencies>Imm32.lib;Kernel32.lib;%(AdditionalDependencies)</AdditionalDependencies>
       <OutputFile>$(SolutionDir)output\$(ProjectName)$(TargetExt)</OutputFile>
-      <UACExecutionLevel>RequireAdministrator</UACExecutionLevel>
+      <UACExecutionLevel>AsInvoker</UACExecutionLevel>
       <LinkTimeCodeGeneration>
       </LinkTimeCodeGeneration>
     </Link>
@@ -155,4 +155,4 @@
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">
   </ImportGroup>
-</Project>
+</Project>
\ No newline at end of file
diff --git a/WeaselSetup/xmake.lua b/WeaselSetup/xmake.lua
index 49439e8bd..865392e46 100644
--- a/WeaselSetup/xmake.lua
+++ b/WeaselSetup/xmake.lua
@@ -4,7 +4,7 @@ target("WeaselSetup")
   add_rules("add_rcfiles", "use_weaselconstants", "subwin")
   add_links("imm32", "kernel32")
 
-  set_policy("windows.manifest.uac", "admin")
+  set_policy("windows.manifest.uac", "invoker")
   add_files("$(projectdir)/PerMonitorHighDPIAware.manifest")
   add_ldflags("/DEBUG /OPT:REF /OPT:ICF /LARGEADDRESSAWARE /ERRORREPORT:QUEUE")