From 4e29475e052c81a68c61931d857872bb9811f270 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sun, 1 Sep 2024 18:47:42 +0800 Subject: [PATCH] fix: bugs in scheduled task (#125) --- Cargo.toml | 2 + src/utils/admin.rs | 20 ++-- src/utils/handle_wrapper.rs | 29 ++++++ src/utils/mod.rs | 2 + src/utils/scheduled_task.rs | 176 ++++++++++++++++++++++++++++++++++-- 5 files changed, 213 insertions(+), 16 deletions(-) create mode 100644 src/utils/handle_wrapper.rs diff --git a/Cargo.toml b/Cargo.toml index 4b020b9..981bd82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,9 +24,11 @@ features = [ "Win32_Graphics_Dwm", "Win32_Graphics_Gdi", "Win32_Security", + "Win32_Security_Authorization", "Win32_System_Console", "Win32_System_LibraryLoader", "Win32_System_Registry", + "Win32_System_SystemInformation", "Win32_System_Threading", "Win32_Storage_FileSystem", ] diff --git a/src/utils/admin.rs b/src/utils/admin.rs index 4e84128..e2f531b 100644 --- a/src/utils/admin.rs +++ b/src/utils/admin.rs @@ -1,6 +1,7 @@ +use super::HandleWrapper; + use anyhow::{anyhow, Result}; use windows::Win32::{ - Foundation::{CloseHandle, HANDLE}, Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION, TOKEN_QUERY}, System::Threading::{GetCurrentProcess, OpenProcessToken}, }; @@ -12,22 +13,23 @@ pub fn is_running_as_admin() -> Result { fn is_running_as_admin_impl() -> Result { let is_elevated = unsafe { - let mut token_handle: HANDLE = HANDLE(0); + let mut token_handle = HandleWrapper::default(); let mut elevation = TOKEN_ELEVATION::default(); let mut returned_length = 0; - OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle)?; + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY, + token_handle.get_handle_mut(), + )?; - let token_information = GetTokenInformation( - token_handle, + GetTokenInformation( + token_handle.get_handle(), TokenElevation, Some(&mut elevation as *mut _ as *mut _), std::mem::size_of::() as u32, &mut returned_length, - ); - - CloseHandle(token_handle)?; + )?; - token_information?; elevation.TokenIsElevated != 0 }; Ok(is_elevated) diff --git a/src/utils/handle_wrapper.rs b/src/utils/handle_wrapper.rs new file mode 100644 index 0000000..0d3a35e --- /dev/null +++ b/src/utils/handle_wrapper.rs @@ -0,0 +1,29 @@ +use windows::Win32::Foundation::{CloseHandle, HANDLE}; + +#[derive(Debug, Clone, Default)] +pub struct HandleWrapper { + handle: HANDLE, +} + +impl HandleWrapper { + pub fn new(handle: HANDLE) -> Self { + Self { handle } + } + pub fn get_handle(&self) -> HANDLE { + self.handle + } + pub fn get_handle_mut(&mut self) -> &mut HANDLE { + &mut self.handle + } +} + +impl Drop for HandleWrapper { + fn drop(&mut self) { + if self.handle.is_invalid() { + return; + } + unsafe { + let _ = CloseHandle(self.handle); + } + } +} diff --git a/src/utils/mod.rs b/src/utils/mod.rs index e81da7d..0fb7867 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,5 +1,6 @@ mod admin; mod check_error; +mod handle_wrapper; mod regedit; mod scheduled_task; mod single_instance; @@ -8,6 +9,7 @@ mod windows_icon; pub use admin::*; pub use check_error::*; +pub use handle_wrapper::*; pub use regedit::*; pub use scheduled_task::*; pub use single_instance::*; diff --git a/src/utils/scheduled_task.rs b/src/utils/scheduled_task.rs index ab55050..e820728 100644 --- a/src/utils/scheduled_task.rs +++ b/src/utils/scheduled_task.rs @@ -1,15 +1,29 @@ -use std::{os::windows::process::CommandExt, process::Command}; +use super::HandleWrapper; -use anyhow::{bail, Result}; -use windows::Win32::System::Threading::CREATE_NO_WINDOW; +use anyhow::{anyhow, bail, Result}; +use std::{ + env, + ffi::OsString, + fs, + os::windows::{ffi::OsStringExt, process::CommandExt}, + process::Command, +}; +use windows::core::{Result as WindowsResult, PWSTR}; +use windows::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER; +use windows::Win32::Security::Authorization::ConvertSidToStringSidW; +use windows::Win32::Security::{ + GetTokenInformation, LookupAccountSidW, TokenUser, SID_NAME_USE, TOKEN_QUERY, TOKEN_USER, +}; +use windows::Win32::System::SystemInformation::GetLocalTime; +use windows::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken, CREATE_NO_WINDOW}; pub fn create_scheduled_task(name: &str, exe_path: &str) -> Result<()> { + let task_xml_path = create_task_file(name, exe_path) + .map_err(|err| anyhow!("Failed to create scheduled task, {err}"))?; + debug!("scheduled task file: {}", task_xml_path); let output = Command::new("schtasks") .creation_flags(CREATE_NO_WINDOW.0) // CREATE_NO_WINDOW flag - .args([ - "/create", "/tn", name, "/tr", exe_path, "/sc", "onlogon", "/rl", "highest", "/it", - "/f", - ]) + .args(["/create", "/tn", name, "/xml", &task_xml_path, "/f"]) .output()?; if !output.status.success() { bail!( @@ -45,3 +59,151 @@ pub fn exist_scheduled_task(name: &str) -> Result { Ok(false) } } + +fn create_task_file(name: &str, exe_path: &str) -> Result { + let (author, user_id) = get_author_and_userid() + .map_err(|err| anyhow!("Failed to get author and user id, {err}"))?; + let current_time = get_current_time(); + let command_path = if exe_path.contains(|c: char| c.is_whitespace()) { + format!("\"{}\"", exe_path) + } else { + exe_path.to_string() + }; + let xml_data = format!( + r#" + + + {current_time} + {author} + \{name} + + + + {current_time} + true + + + + + {user_id} + InteractiveToken + HighestAvailable + + + + IgnoreNew + false + true + true + false + false + + true + false + + true + true + false + false + false + PT0S + 7 + + + + {command_path} + + +"# + ); + let xml_path = env::temp_dir().join("window-switcher-task.xml"); + let xml_path = xml_path.display().to_string(); + fs::write(&xml_path, xml_data) + .map_err(|err| anyhow!("Failed to write task xml file at '{xml_path}', {err}",))?; + Ok(xml_path) +} + +fn get_author_and_userid() -> WindowsResult<(String, String)> { + let mut token_handle = HandleWrapper::default(); + unsafe { + OpenProcessToken( + GetCurrentProcess(), + TOKEN_QUERY, + token_handle.get_handle_mut(), + )? + }; + + let mut token_info_length = 0; + if let Err(err) = unsafe { + GetTokenInformation( + token_handle.get_handle(), + TokenUser, + None, + 0, + &mut token_info_length, + ) + } { + if err != ERROR_INSUFFICIENT_BUFFER.into() { + return Err(err); + } + } + + let mut token_user = Vec::::with_capacity(token_info_length as usize); + unsafe { + GetTokenInformation( + token_handle.get_handle(), + TokenUser, + Some(token_user.as_mut_ptr() as *mut _), + token_info_length, + &mut token_info_length, + )? + }; + + let user_sid = unsafe { *(token_user.as_ptr() as *const TOKEN_USER) } + .User + .Sid; + + let mut name = Vec::::with_capacity(256); + let mut name_len = 256; + let mut domain = Vec::::with_capacity(256); + let mut domain_len = 256; + let mut sid_name_use = SID_NAME_USE(0); + + unsafe { + LookupAccountSidW( + None, + user_sid, + PWSTR(name.as_mut_ptr()), + &mut name_len, + PWSTR(domain.as_mut_ptr()), + &mut domain_len, + &mut sid_name_use, + )? + }; + + unsafe { + name.set_len(name_len as usize); + domain.set_len(domain_len as usize); + } + + let username = OsString::from_wide(&name).to_string_lossy().into_owned(); + let domainname = OsString::from_wide(&domain).to_string_lossy().into_owned(); + + let mut sid_string = PWSTR::null(); + unsafe { ConvertSidToStringSidW(user_sid, &mut sid_string)? }; + + let sid_str = OsString::from_wide(unsafe { sid_string.as_wide() }) + .to_string_lossy() + .into_owned(); + + Ok((format!("{}\\{}", domainname, username), sid_str)) +} + +fn get_current_time() -> String { + let st = unsafe { GetLocalTime() }; + + format!( + "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}", + st.wYear, st.wMonth, st.wDay, st.wHour, st.wMinute, st.wSecond, + ) +}