diff --git a/Cargo.toml b/Cargo.toml index dd8fd8b..be6f7cd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,4 +10,4 @@ subprocess = "0.2.9" yaml-rust = "0.4.5" [target.'cfg(windows)'.dependencies] -winapi = { version = "0.3", features = ["errhandlingapi", "impl-default", "ioapiset", "jobapi2", "wincon", "winuser"] } +winapi = { version = "0.3", features = ["errhandlingapi", "impl-default", "ioapiset", "jobapi2", "memoryapi", "wincon", "winuser"] } diff --git a/src/cfg.rs b/src/cfg.rs index 81ddfe0..6636bb3 100644 --- a/src/cfg.rs +++ b/src/cfg.rs @@ -157,4 +157,9 @@ impl Config { .map(|s| s.to_owned()) .unwrap_or(false) } + + #[cfg(windows)] + pub fn hook_dll(&self) -> Vec { + self.get_str_vec("hook_dll").unwrap_or(vec![]) + } } diff --git a/src/main.rs b/src/main.rs index 14a4396..99bacc3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -116,7 +116,13 @@ impl Main { #[cfg(windows)] fn call(cml: Vec) -> Result { - windows::call(&cml).map(|c| ExitStatus::Exited(c)) + let t = Vec::::new(); + windows::call(&cml, &t).map(|c| ExitStatus::Exited(c)) + } + + #[cfg(windows)] + fn call2(cml: Vec, dlls: Vec) -> Result { + windows::call(&cml, &dlls).map(|c| ExitStatus::Exited(c)) } fn restore(&self) -> Result<(), Error> { @@ -210,8 +216,11 @@ impl Main { if need_hide && !hide { println!("Failed to hide console window."); } + #[cfg(not(windows))] let e = Self::call(cml)?; #[cfg(windows)] + let e = Self::call2(cml, self._cfg.hook_dll())?; + #[cfg(windows)] if hide { windows::show_window(); } diff --git a/src/windows.rs b/src/windows.rs index 3843256..defee54 100644 --- a/src/windows.rs +++ b/src/windows.rs @@ -8,16 +8,20 @@ use winapi::shared::minwindef::DWORD; use winapi::um::errhandlingapi::GetLastError; use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE}; use winapi::um::ioapiset::{CreateIoCompletionPort, GetQueuedCompletionStatus}; +use winapi::um::libloaderapi::{GetModuleHandleA, GetProcAddress}; use winapi::um::jobapi2::{AssignProcessToJobObject, SetInformationJobObject}; +use winapi::um::memoryapi::{VirtualAllocEx, WriteProcessMemory, VirtualFreeEx}; use winapi::um::minwinbase::LPOVERLAPPED; use winapi::um::processthreadsapi::{ - CreateProcessW, GetExitCodeProcess, ResumeThread, PROCESS_INFORMATION, STARTUPINFOW, + CreateProcessW, GetExitCodeProcess, ResumeThread, TerminateProcess, PROCESS_INFORMATION, + STARTUPINFOW, CreateRemoteThread, }; +use winapi::um::synchapi::WaitForSingleObject; use winapi::um::winbase::{CreateJobObjectA, CREATE_SUSPENDED, INFINITE}; use winapi::um::wincon::GetConsoleWindow; use winapi::um::winnt::{ JobObjectAssociateCompletionPortInformation, JOBOBJECT_ASSOCIATE_COMPLETION_PORT, - JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO, + JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO, MEM_COMMIT, PAGE_READWRITE, MEM_RELEASE, }; use winapi::um::winuser::{ShowWindow, SW_HIDE, SW_SHOW}; @@ -43,9 +47,10 @@ pub enum PopenError { CreateJobFailed, CreateProcessFailed, AssignJobFailed, + CreateThreadFailed, } -pub fn call>(argv: &[S]) -> Result { +pub fn call, T: AsRef>(argv: &[S], dlls: &[T]) -> Result { let job = unsafe { CreateJobObjectA(null_mut(), null()) }; if job.is_null() { println!("Failed to create job: {}.", unsafe { GetLastError() }); @@ -129,6 +134,71 @@ pub fn call>(argv: &[S]) -> Result { unsafe { CloseHandle(io_port) }; return Err(PopenError::AssignJobFailed); } + for i in dlls.iter() { + let dll: Vec<_> = i.as_ref().encode_wide().collect(); + let mem_size = (dll.len() + 1) * size_of::(); + let p_dll_path = unsafe { + VirtualAllocEx( + pi.hProcess, + null_mut(), + mem_size, + MEM_COMMIT, + PAGE_READWRITE, + ) + }; + if p_dll_path.is_null() { + println!("Failed to allocate memory in remote process."); + unsafe { TerminateProcess(pi.hProcess, 1) }; + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(io_port) }; + unsafe { CloseHandle(pi.hProcess) }; + unsafe { CloseHandle(pi.hThread) }; + return Err(PopenError::CreateProcessFailed); + } + let re = unsafe { + WriteProcessMemory( + pi.hProcess, + p_dll_path, + dll.as_ptr() as *const c_void, + mem_size, + null_mut(), + ) != 0 + }; + if !re { + println!("Failed to write memory in remote process."); + unsafe { VirtualFreeEx(pi.hProcess, p_dll_path, 0, MEM_RELEASE) }; + unsafe { TerminateProcess(pi.hProcess, 1) }; + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(io_port) }; + unsafe { CloseHandle(pi.hProcess) }; + unsafe { CloseHandle(pi.hThread) }; + return Err(PopenError::CreateProcessFailed); + } + let h_thread = unsafe { + CreateRemoteThread( + pi.hProcess, + null_mut(), + 0, + Some(std::mem::transmute(GetProcAddress(GetModuleHandleA("kernel32\0".as_ptr() as *const i8), "LoadLibraryW\0".as_ptr() as *const i8))), + p_dll_path, + 0, + null_mut(), + ) + }; + if h_thread.is_null() { + println!("Failed to create remote thread."); + unsafe { VirtualFreeEx(pi.hProcess, p_dll_path, 0, MEM_RELEASE) }; + unsafe { TerminateProcess(pi.hProcess, 1) }; + unsafe { CloseHandle(job) }; + unsafe { CloseHandle(io_port) }; + unsafe { CloseHandle(pi.hProcess) }; + unsafe { CloseHandle(pi.hThread) }; + return Err(PopenError::CreateThreadFailed); + } + unsafe { WaitForSingleObject(h_thread, INFINITE) }; + unsafe { VirtualFreeEx(pi.hProcess, p_dll_path, 0, MEM_RELEASE) }; + unsafe { CloseHandle(h_thread) }; + } unsafe { ResumeThread(pi.hThread) }; let mut code = DWORD::default(); let mut key = ULONG_PTR::default();