helios-remote/crates/client/src/main.rs

757 lines
30 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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::<u32>() {
// 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<Self> {
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<String> = 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>) -> String;
fn unwrap_or_else(self, f: impl FnOnce() -> String) -> String;
}
impl OrElseString for String {
fn or_else(self, f: impl FnOnce() -> Option<String>) -> 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<Mutex<shell::PersistentShell>>,
) -> 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<u16> = msg.encode_utf16().chain(std::iter::once(0)).collect();
let ttl_w: Vec<u16> = 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<String> = 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() }
}
}
}
}