use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use colored::Colorize; use futures_util::{SinkExt, StreamExt}; use native_tls::TlsConnector; use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; use tokio_tungstenite::{connect_async_tls_with_config, tungstenite::Message, Connector}; use base64::Engine; use helios_common::{ClientMessage, ServerMessage}; #[allow(unused_imports)] use reqwest; use helios_common::protocol::{is_valid_label, sanitize_label}; mod display; mod logger; mod shell; mod screenshot; mod windows_mgmt; use display::trunc; fn banner() { println!(); println!(" {} {}", "☀ ".yellow().bold(), "HELIOS REMOTE".bold()); display::info_line("🔗", "commit:", &env!("GIT_COMMIT").dimmed().to_string()); } fn print_device_info(label: &str) { #[cfg(windows)] { let admin = is_admin(); let priv_str = if admin { "admin".dimmed().to_string() } else { "no admin".dimmed().to_string() }; display::info_line("👤", "privileges:", &priv_str); } #[cfg(not(windows))] display::info_line("👤", "privileges:", &"no admin".dimmed().to_string()); display::info_line("🖥", "device:", &label.dimmed().to_string()); println!(); } #[cfg(windows)] fn is_admin() -> bool { use windows::Win32::UI::Shell::IsUserAnAdmin; unsafe { IsUserAnAdmin().as_bool() } } #[cfg(windows)] fn enable_ansi() { use windows::Win32::System::Console::{ GetConsoleMode, GetStdHandle, SetConsoleMode, ENABLE_VIRTUAL_TERMINAL_PROCESSING, STD_OUTPUT_HANDLE, }; unsafe { if let Ok(handle) = GetStdHandle(STD_OUTPUT_HANDLE) { let mut mode = Default::default(); if GetConsoleMode(handle, &mut mode).is_ok() { let _ = SetConsoleMode(handle, mode | ENABLE_VIRTUAL_TERMINAL_PROCESSING); } } } } // ── Single instance enforcement ───────────────────────────────────────────── fn lock_file_path() -> PathBuf { let base = dirs::config_dir() .or_else(|| dirs::home_dir()) .unwrap_or_else(|| PathBuf::from(".")); base.join("helios-remote").join("instance.lock") } /// Try to acquire a single-instance lock. Returns true if we got it. fn acquire_instance_lock() -> bool { let path = lock_file_path(); if let Some(parent) = path.parent() { let _ = std::fs::create_dir_all(parent); } // Check if another instance is running if path.exists() { if let Ok(content) = std::fs::read_to_string(&path) { if let Ok(pid) = content.trim().parse::() { // Check if process is still alive #[cfg(windows)] { use windows::Win32::System::Threading::{OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION}; let alive = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, pid).is_ok() }; if alive { return false; } } #[cfg(not(windows))] { use std::process::Command; let alive = Command::new("kill") .args(["-0", &pid.to_string()]) .status() .map(|s| s.success()) .unwrap_or(false); if alive { return false; } } } } } // Write our PID let pid = std::process::id(); std::fs::write(&path, pid.to_string()).is_ok() } fn release_instance_lock() { let _ = std::fs::remove_file(lock_file_path()); } // ── Config ────────────────────────────────────────────────────────────────── #[derive(Debug, Clone, Serialize, Deserialize)] struct Config { relay_url: String, api_key: String, label: String, } impl Config { fn config_path() -> PathBuf { let base = dirs::config_dir() .or_else(|| dirs::home_dir()) .unwrap_or_else(|| PathBuf::from(".")); base.join("helios-remote").join("config.toml") } fn load() -> Option { let path = Self::config_path(); let data = std::fs::read_to_string(&path).ok()?; toml::from_str(&data).ok() } fn save(&self) -> std::io::Result<()> { let path = Self::config_path(); std::fs::create_dir_all(path.parent().unwrap())?; let data = toml::to_string_pretty(self).unwrap(); std::fs::write(&path, data)?; Ok(()) } } fn prompt_config() -> Config { use std::io::Write; let relay_url = { let default = "wss://remote.agent-helios.me/ws"; print!(" {} Relay URL [{}]: ", "→".cyan().bold(), default); std::io::stdout().flush().unwrap(); let mut input = String::new(); std::io::stdin().read_line(&mut input).unwrap(); let trimmed = input.trim(); if trimmed.is_empty() { default.to_string() } else { trimmed.to_string() } }; let api_key = { print!(" {} API Key: ", "→".cyan().bold()); std::io::stdout().flush().unwrap(); let mut input = String::new(); std::io::stdin().read_line(&mut input).unwrap(); input.trim().to_string() }; let label = { let default_label = sanitize_label(&hostname()); loop { print!(" {} Device label [{}]: ", "→".cyan().bold(), default_label); std::io::stdout().flush().unwrap(); let mut input = String::new(); std::io::stdin().read_line(&mut input).unwrap(); let trimmed = input.trim(); let candidate = if trimmed.is_empty() { default_label.clone() } else { trimmed.to_string() }; if is_valid_label(&candidate) { break candidate; } println!(" {} Label must be lowercase, no spaces. Only a-z, 0-9, '-', '_'.", "✗".red().bold()); println!(" Suggestion: {}", sanitize_label(&candidate).cyan()); } }; Config { relay_url, api_key, label } } #[tokio::main] async fn main() { #[cfg(windows)] enable_ansi(); logger::init(); if std::env::var("RUST_LOG").is_err() { unsafe { std::env::set_var("RUST_LOG", "off"); } } banner(); // Clean up leftover .old.exe from previous self-update (Windows can't delete running exe) #[cfg(target_os = "windows")] if let Ok(exe) = std::env::current_exe() { let old = exe.with_extension("old.exe"); let _ = std::fs::remove_file(&old); } // Single instance check if !acquire_instance_lock() { display::err("❌", "Another instance of helios-remote is already running."); display::err("", "Only one instance per device is allowed."); std::process::exit(1); } // Clean up lock on exit let _guard = scopeguard::guard((), |_| release_instance_lock()); // Load or prompt for config let config = match Config::load() { Some(c) => { // Validate existing label if !is_valid_label(&c.label) { let new_label = sanitize_label(&c.label); display::info_line("⚠", "migrate:", &format!( "Label '{}' is invalid, migrating to '{}'", c.label, new_label )); let mut cfg = c; cfg.label = new_label; if let Err(e) = cfg.save() { display::err("❌", &format!("Failed to save config: {e}")); } cfg } else { c } } None => { display::info_line("ℹ", "setup:", "No config found — first-time setup"); println!(); let c = prompt_config(); println!(); if let Err(e) = c.save() { display::err("❌", &format!("Failed to save config: {e}")); } else { display::info_line("✅", "config:", "saved"); } // Self-restart after first-time setup so all config takes effect cleanly println!(); display::info_line("🔄", "restart:", "Config saved. Restarting..."); release_instance_lock(); let exe = std::env::current_exe().expect("Failed to get current exe path"); let args: Vec = std::env::args().skip(1).collect(); let _ = std::process::Command::new(exe).args(&args).spawn(); std::process::exit(0); } }; let label = config.label.clone(); print_device_info(&label); let config = Arc::new(config); let shell = Arc::new(Mutex::new(shell::PersistentShell::new())); // Connect with exponential backoff let mut backoff = Duration::from_secs(1); const MAX_BACKOFF: Duration = Duration::from_secs(30); loop { let host = config.relay_url .trim_start_matches("wss://") .trim_start_matches("ws://") .split('/') .next() .unwrap_or(&config.relay_url); display::cmd_start("🌐", "connect", host); let tls_connector = TlsConnector::builder() .danger_accept_invalid_certs(true) .build() .expect("TLS connector build failed"); let connector = Connector::NativeTls(tls_connector); match connect_async_tls_with_config(&config.relay_url, None, false, Some(connector)).await { Ok((ws_stream, _)) => { display::cmd_done("🌐", "connect", host, true, "connected"); backoff = Duration::from_secs(1); let (mut write, mut read) = ws_stream.split(); // Send Hello with device label let hello = ClientMessage::Hello { label: label.clone(), }; let hello_json = serde_json::to_string(&hello).unwrap(); if let Err(e) = write.send(Message::Text(hello_json)).await { display::err("❌", &format!("hello failed: {e}")); tokio::time::sleep(backoff).await; backoff = (backoff * 2).min(MAX_BACKOFF); continue; } let write = Arc::new(Mutex::new(write)); while let Some(msg_result) = read.next().await { match msg_result { Ok(Message::Text(text)) => { let server_msg: ServerMessage = match serde_json::from_str(&text) { Ok(m) => m, Err(e) => { display::err("❌", &format!("Failed to parse server message: {e}")); continue; } }; let write_clone = Arc::clone(&write); let shell_clone = Arc::clone(&shell); tokio::spawn(async move { let response = handle_message(server_msg, shell_clone).await; let json = match serde_json::to_string(&response) { Ok(j) => j, Err(e) => { display::err("❌", &format!("Failed to serialize response: {e}")); return; } }; let mut w = write_clone.lock().await; if let Err(e) = w.send(Message::Text(json)).await { display::err("❌", &format!("Failed to send response: {e}")); } }); } Ok(Message::Ping(data)) => { let mut w = write.lock().await; let _ = w.send(Message::Pong(data)).await; } Ok(Message::Close(_)) => { display::cmd_start("🌐", "connect", host); display::cmd_done("🌐", "connect", host, false, "connection lost"); break; } Err(e) => { display::cmd_done("🌐", "connect", host, false, &format!("lost: {e}")); break; } _ => {} } } } Err(e) => { display::cmd_start("🌐", "connect", host); display::cmd_done("🌐", "connect", host, false, &format!("{e}")); } } tokio::time::sleep(backoff).await; backoff = (backoff * 2).min(MAX_BACKOFF); } } fn hostname() -> String { std::fs::read_to_string("/etc/hostname") .unwrap_or_default() .trim() .to_string() .or_else(|| std::env::var("COMPUTERNAME").ok()) .unwrap_or_else(|| "unknown".to_string()) } trait OrElseString { fn or_else(self, f: impl FnOnce() -> Option) -> String; fn unwrap_or_else(self, f: impl FnOnce() -> String) -> String; } impl OrElseString for String { fn or_else(self, f: impl FnOnce() -> Option) -> String { if self.is_empty() { f().unwrap_or_default() } else { self } } fn unwrap_or_else(self, f: impl FnOnce() -> String) -> String { if self.is_empty() { f() } else { self } } } async fn handle_message( msg: ServerMessage, shell: Arc>, ) -> ClientMessage { match msg { ServerMessage::WindowScreenshotRequest { request_id, window_id } => { let payload = format!("window {window_id}"); display::cmd_start("📷", "screenshot", &payload); match screenshot::take_window_screenshot(window_id) { Ok((image_base64, width, height)) => { display::cmd_done("📷", "screenshot", &payload, true, &format!("{width}×{height}")); ClientMessage::ScreenshotResponse { request_id, image_base64, width, height } } Err(e) => { display::cmd_done("📷", "screenshot", &payload, false, &format!("{e}")); ClientMessage::Error { request_id, message: format!("Window screenshot failed: {e}") } } } } ServerMessage::ScreenshotRequest { request_id } => { display::cmd_start("📷", "screenshot", "screen"); match screenshot::take_screenshot() { Ok((image_base64, width, height)) => { display::cmd_done("📷", "screenshot", "screen", true, &format!("{width}×{height}")); ClientMessage::ScreenshotResponse { request_id, image_base64, width, height } } Err(e) => { display::cmd_done("📷", "screenshot", "screen", false, &format!("{e}")); ClientMessage::Error { request_id, message: format!("Screenshot failed: {e}") } } } } ServerMessage::InformRequest { request_id, message, title } => { let msg = message.clone(); let ttl = title.clone().unwrap_or_else(|| "Helios".to_string()); // Fire-and-forget: show MessageBox in background thread, don't block std::thread::spawn(move || { #[cfg(windows)] unsafe { use windows::core::PCWSTR; use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION, HWND_DESKTOP}; let msg_w: Vec = msg.encode_utf16().chain(std::iter::once(0)).collect(); let ttl_w: Vec = ttl.encode_utf16().chain(std::iter::once(0)).collect(); MessageBoxW(HWND_DESKTOP, PCWSTR(msg_w.as_ptr()), PCWSTR(ttl_w.as_ptr()), MB_OK | MB_ICONINFORMATION); } #[cfg(not(windows))] let _ = (msg, ttl); }); display::cmd_done("📢", "inform", &message, true, "sent"); ClientMessage::Ack { request_id } } ServerMessage::PromptRequest { request_id, message, title: _ } => { display::prompt_waiting(&message); let answer = tokio::task::spawn_blocking(|| { let mut input = String::new(); std::io::stdin().read_line(&mut input).ok(); input.trim().to_string() }).await.unwrap_or_default(); display::prompt_done(&message, &answer); ClientMessage::PromptResponse { request_id, answer } } ServerMessage::ExecRequest { request_id, command, timeout_ms } => { let payload = trunc(&command, 80); display::cmd_start("⚡", "execute", &payload); let mut sh = shell.lock().await; match sh.run(&command, timeout_ms).await { Ok((stdout, stderr, exit_code)) => { let result = if exit_code != 0 { let err_line = stderr.lines() .map(|l| l.trim()) .find(|l| !l.is_empty() && !l.starts_with("In Zeile:") && !l.starts_with("+ ") && !l.starts_with("CategoryInfo") && !l.starts_with("FullyQualifiedErrorId")) .unwrap_or("error") .to_string(); err_line } else { stdout.trim().lines().next().unwrap_or("").to_string() }; display::cmd_done("⚡", "execute", &payload, exit_code == 0, &result); ClientMessage::ExecResponse { request_id, stdout, stderr, exit_code } } Err(e) => { display::cmd_done("⚡", "execute", &payload, false, &format!("exec failed: {e}")); ClientMessage::Error { request_id, message: format!("Exec failed for command {:?}.\nError: {e}", command) } } } } ServerMessage::ListWindowsRequest { request_id } => { display::cmd_start("🪟", "list windows", ""); match windows_mgmt::list_windows() { Ok(windows) => { display::cmd_done("🪟", "list windows", "", true, &format!("{} windows", windows.len())); ClientMessage::ListWindowsResponse { request_id, windows } } Err(e) => { display::cmd_done("🪟", "list windows", "", false, &e); ClientMessage::Error { request_id, message: e } } } } ServerMessage::MinimizeAllRequest { request_id } => { display::cmd_start("🪟", "minimize all", ""); match windows_mgmt::minimize_all() { Ok(()) => { display::cmd_done("🪟", "minimize all", "", true, "done"); ClientMessage::Ack { request_id } } Err(e) => { display::cmd_done("🪟", "minimize all", "", false, &e); ClientMessage::Error { request_id, message: e } } } } ServerMessage::FocusWindowRequest { request_id, window_id } => { let payload = format!("{window_id}"); display::cmd_start("🪟", "focus window", &payload); match windows_mgmt::focus_window(window_id) { Ok(()) => { display::cmd_done("🪟", "focus window", &payload, true, "done"); ClientMessage::Ack { request_id } } Err(e) => { display::cmd_done("🪟", "focus window", &payload, false, &e); ClientMessage::Error { request_id, message: e } } } } ServerMessage::MaximizeAndFocusRequest { request_id, window_id } => { let payload = format!("{window_id}"); display::cmd_start("🪟", "maximize", &payload); match windows_mgmt::maximize_and_focus(window_id) { Ok(()) => { display::cmd_done("🪟", "maximize", &payload, true, "done"); ClientMessage::Ack { request_id } } Err(e) => { display::cmd_done("🪟", "maximize", &payload, false, &e); ClientMessage::Error { request_id, message: e } } } } ServerMessage::VersionRequest { request_id } => { display::cmd_start("ℹ", "version", ""); let version = env!("CARGO_PKG_VERSION").to_string(); let commit = env!("GIT_COMMIT").to_string(); display::cmd_done("ℹ", "version", "", true, &commit); ClientMessage::VersionResponse { request_id, version, commit } } ServerMessage::LogsRequest { request_id, lines } => { let payload = format!("last {lines} lines"); display::cmd_start("📜", "logs", &payload); let content = logger::tail(lines); let log_path = logger::get_log_path(); display::cmd_done("📜", "logs", &payload, true, &log_path); ClientMessage::LogsResponse { request_id, content, log_path } } ServerMessage::UploadRequest { request_id, path, content_base64 } => { let payload = trunc(&path, 60); display::cmd_start("📁", "upload", &payload); match (|| -> Result<(), String> { let bytes = base64::engine::general_purpose::STANDARD .decode(&content_base64) .map_err(|e| format!("base64 decode: {e}"))?; if let Some(parent) = std::path::Path::new(&path).parent() { std::fs::create_dir_all(parent).map_err(|e| e.to_string())?; } std::fs::write(&path, &bytes).map_err(|e| e.to_string())?; Ok(()) })() { Ok(()) => { display::cmd_done("📁", "upload", &payload, true, "saved"); ClientMessage::Ack { request_id } } Err(e) => { display::cmd_done("📁", "upload", &payload, false, &e); ClientMessage::Error { request_id, message: e } } } } ServerMessage::DownloadRequest { request_id, path } => { let payload = trunc(&path, 60); display::cmd_start("📁", "download", &payload); match std::fs::read(&path) { Ok(bytes) => { let size = bytes.len() as u64; let content_base64 = base64::engine::general_purpose::STANDARD.encode(&bytes); display::cmd_done("📁", "download", &payload, true, &format!("{size} bytes")); ClientMessage::DownloadResponse { request_id, content_base64, size } } Err(e) => { display::cmd_done("📁", "download", &payload, false, &format!("read failed: {e}")); ClientMessage::Error { request_id, message: format!("Read failed: {e}") } } } } ServerMessage::RunRequest { request_id, program, args } => { let payload = if args.is_empty() { program.clone() } else { format!("{program} {}", args.join(" ")) }; let payload = trunc(&payload, 60); display::cmd_start("🚀", "run", &payload); use std::process::Command as StdCommand; match StdCommand::new(&program).args(&args).spawn() { Ok(_) => { display::cmd_done("🚀", "run", &payload, true, "started"); ClientMessage::Ack { request_id } } Err(e) => { display::cmd_done("🚀", "run", &payload, false, &format!("{e}")); ClientMessage::Error { request_id, message: format!("Failed to start '{}': {e}", program) } } } } ServerMessage::ClipboardGetRequest { request_id } => { display::cmd_start("📋", "get clipboard", ""); let out = tokio::process::Command::new("powershell.exe") .args(["-NoProfile", "-NonInteractive", "-Command", "Get-Clipboard"]) .output().await; match out { Ok(o) => { let text = String::from_utf8_lossy(&o.stdout).trim().to_string(); display::cmd_done("📋", "get clipboard", "", true, &text); ClientMessage::ClipboardGetResponse { request_id, text } } Err(e) => { display::cmd_done("📋", "get clipboard", "", false, &format!("{e}")); ClientMessage::Error { request_id, message: format!("Clipboard get failed: {e}") } } } } ServerMessage::ClipboardSetRequest { request_id, text } => { let payload = trunc(&text, 60); display::cmd_start("📋", "set clipboard", &payload); let cmd = format!("Set-Clipboard -Value '{}'", text.replace('\'', "''")); let out = tokio::process::Command::new("powershell.exe") .args(["-NoProfile", "-NonInteractive", "-Command", &cmd]) .output().await; match out { Ok(_) => { display::cmd_done("📋", "set clipboard", &payload, true, &payload); ClientMessage::Ack { request_id } } Err(e) => { display::cmd_done("📋", "set clipboard", &payload, false, &format!("{e}")); ClientMessage::Error { request_id, message: format!("Clipboard set failed: {e}") } } } } ServerMessage::UpdateRequest { request_id } => { display::cmd_start("🔄", "update", "downloading..."); let exe = std::env::current_exe().ok(); tokio::spawn(async move { // Give the response time to be sent before we restart tokio::time::sleep(tokio::time::Duration::from_millis(800)).await; let exe = match exe { Some(p) => p, None => { display::err("❌", "update: could not determine current exe path"); return; } }; let url = "https://agent-helios.me/downloads/helios-remote/helios-remote-client-windows.exe"; let bytes = match reqwest::get(url).await { Ok(r) => match r.bytes().await { Ok(b) => b, Err(e) => { display::err("❌", &format!("update: read body failed: {e}")); return; } }, Err(e) => { display::err("❌", &format!("update: download failed: {e}")); return; } }; // Write new binary to a temp path, then swap let tmp = exe.with_extension("update.exe"); if let Err(e) = std::fs::write(&tmp, &bytes) { display::err("❌", &format!("update: write failed: {e}")); return; } // Rename current → .old, then tmp → current let old = exe.with_extension("old.exe"); let _ = std::fs::remove_file(&old); if let Err(e) = std::fs::rename(&exe, &old) { display::err("❌", &format!("update: rename old failed: {e}")); return; } if let Err(e) = std::fs::rename(&tmp, &exe) { // Attempt rollback let _ = std::fs::rename(&old, &exe); display::err("❌", &format!("update: rename new failed: {e}")); return; } display::cmd_done("🔄", "update", "", true, "updated — restarting"); // Delete old binary let _ = std::fs::remove_file(&old); // Restart with same args (new console window on Windows) let args: Vec = std::env::args().skip(1).collect(); #[cfg(target_os = "windows")] { // Use "start" to open a new visible console window let exe_str = exe.to_string_lossy(); let _ = std::process::Command::new("cmd") .args(["/c", "start", "", &exe_str]) .spawn(); } #[cfg(not(target_os = "windows"))] let _ = std::process::Command::new(&exe).args(&args).spawn(); std::process::exit(0); }); display::cmd_done("🔄", "update", "", true, "triggered"); ClientMessage::UpdateResponse { request_id, success: true, message: "update triggered, client restarting...".into(), } } ServerMessage::Ack { request_id } => { ClientMessage::Ack { request_id } } ServerMessage::Error { request_id, message } => { display::err("❌", &format!("server error: {message}")); if let Some(rid) = request_id { ClientMessage::Ack { request_id: rid } } else { ClientMessage::Hello { label: String::new() } } } } }