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

673 lines
24 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};
use uuid::Uuid;
mod logger;
mod shell;
mod screenshot;
mod input;
mod windows_mgmt;
// ─── CLI Output ─────────────────────────────────────────────────────────────
/// Truncate a string to at most `max_chars` Unicode characters.
fn trunc(s: &str, max_chars: usize) -> String {
let mut chars = s.chars();
let truncated: String = chars.by_ref().take(max_chars).collect();
if chars.next().is_some() {
format!("{}", truncated)
} else {
truncated
}
}
fn banner() {
println!();
println!(" {}\tHELIOS REMOTE ({})", "".yellow().bold(), env!("GIT_COMMIT"));
#[cfg(windows)]
{
let admin = is_admin();
let (icon, admin_str) = if admin {
("👑", "admin".green().bold().to_string())
} else {
("👤", "user (no admin)".yellow().to_string())
};
println!(" {}\t{}", icon, admin_str);
}
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);
}
}
}
}
macro_rules! log_status {
($($arg:tt)*) => {
println!(" {}", format!($($arg)*));
};
}
macro_rules! log_ok {
($($arg:tt)*) => {{
let msg = format!($($arg)*);
println!(" {}\t{}", "", msg);
logger::write_line("OK", &msg);
}};
}
macro_rules! log_err {
($($arg:tt)*) => {{
let msg = format!($($arg)*);
println!(" {}\t{}", "", msg);
logger::write_line("ERROR", &msg);
}};
}
macro_rules! log_cmd {
($emoji:expr, $($arg:tt)*) => {{
let msg = format!($($arg)*);
println!(" {}\t{}", $emoji, msg);
logger::write_line("CMD", &msg);
}};
}
// ────────────────────────────────────────────────────────────────────────────
#[derive(Debug, Clone, Serialize, Deserialize)]
struct Config {
relay_url: String,
api_key: String,
label: Option<String>,
session_id: Option<String>, // persistent UUID
}
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 = hostname();
print!(" {} Label for this PC [{}]: ", "".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().to_string();
if trimmed.is_empty() {
Some(default_label)
} else {
Some(trimmed)
}
};
Config { relay_url, api_key, label, session_id: None }
}
#[tokio::main]
async fn main() {
// Enable ANSI color codes on Windows (required when running as admin)
#[cfg(windows)]
enable_ansi();
logger::init();
// Suppress tracing output by default
if std::env::var("RUST_LOG").is_err() {
unsafe { std::env::set_var("RUST_LOG", "off"); }
}
banner();
// Load or prompt for config
let config = match Config::load() {
Some(c) => c,
None => {
log_status!("No config found — first-time setup");
println!();
let c = prompt_config();
println!();
if let Err(e) = c.save() {
log_err!("Failed to save config: {e}");
} else {
log_ok!("Config saved");
}
c
}
};
// Resolve or generate persistent session UUID
let sid: Uuid = match &config.session_id {
Some(id) => Uuid::parse_str(id).unwrap_or_else(|_| Uuid::new_v4()),
None => {
let id = Uuid::new_v4();
let mut cfg = config.clone();
cfg.session_id = Some(id.to_string());
if let Err(e) = cfg.save() {
log_err!("Failed to save session_id: {e}");
}
id
}
};
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);
log_status!("🌐 Connecting to {}...", 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, _)) => {
let label = config.label.clone().unwrap_or_else(|| hostname());
log_ok!(
"Connected {} {} {} Session {}",
"·".dimmed(),
label.bold(),
"·".dimmed(),
sid.to_string().dimmed()
);
println!();
backoff = Duration::from_secs(1);
let (mut write, mut read) = ws_stream.split();
// Send Hello
let hello = ClientMessage::Hello {
label: config.label.clone(),
};
let hello_json = serde_json::to_string(&hello).unwrap();
if let Err(e) = write.send(Message::Text(hello_json)).await {
log_err!("Failed to send Hello: {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) => {
log_err!("Failed to parse server message: {e}");
continue;
}
};
let write_clone = Arc::clone(&write);
let shell_clone = Arc::clone(&shell);
tokio::spawn(async move {
// tokio isolates panics per task — a panic here won't kill
// the main loop. handle_message uses map_err everywhere so
// it should never panic in practice.
let response = handle_message(server_msg, shell_clone).await;
let json = match serde_json::to_string(&response) {
Ok(j) => j,
Err(e) => {
log_err!("Failed to serialize response: {e}");
return;
}
};
let mut w = write_clone.lock().await;
if let Err(e) = w.send(Message::Text(json)).await {
log_err!("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(_)) => {
log_err!("Connection lost — reconnecting...");
break;
}
Err(e) => {
log_err!("Connection lost: {e} — reconnecting...");
break;
}
_ => {}
}
}
}
Err(e) => {
log_err!("Connection failed: {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 } => {
log_cmd!("📷", "screenshot window {}", window_id);
match screenshot::take_window_screenshot(window_id) {
Ok((image_base64, width, height)) => {
log_ok!("Done {} {}×{}", "·".dimmed(), width, height);
ClientMessage::ScreenshotResponse { request_id, image_base64, width, height }
}
Err(e) => {
log_err!("Window screenshot failed: {e}");
ClientMessage::Error { request_id, message: format!("Window screenshot failed: {e}") }
}
}
}
ServerMessage::ScreenshotRequest { request_id } => {
log_cmd!("📷", "screenshot");
match screenshot::take_screenshot() {
Ok((image_base64, width, height)) => {
log_ok!("Done {} {}×{}", "·".dimmed(), width, height);
ClientMessage::ScreenshotResponse {
request_id,
image_base64,
width,
height,
}
}
Err(e) => {
log_err!("Screenshot failed: {e}");
ClientMessage::Error {
request_id,
message: format!("Screenshot failed: {e}"),
}
}
}
}
ServerMessage::PromptRequest { request_id, message, title } => {
let _title = title.unwrap_or_else(|| "Helios Remote".to_string());
#[cfg(windows)]
let title = _title.clone();
log_cmd!("💬", "prompt {}", trunc(&message, 60));
#[cfg(windows)]
{
use windows::core::PCWSTR;
use windows::Win32::UI::WindowsAndMessaging::{MessageBoxW, MB_OK, MB_ICONINFORMATION, HWND_DESKTOP};
let msg_wide: Vec<u16> = message.encode_utf16().chain(std::iter::once(0)).collect();
let title_wide: Vec<u16> = title.encode_utf16().chain(std::iter::once(0)).collect();
// Run blocking MessageBox in a thread so we don't block the async runtime
let msg_clone = message.clone();
tokio::task::spawn_blocking(move || {
unsafe {
MessageBoxW(
HWND_DESKTOP,
PCWSTR(msg_wide.as_ptr()),
PCWSTR(title_wide.as_ptr()),
MB_OK | MB_ICONINFORMATION,
);
}
}).await.ok();
log_ok!("User confirmed: {}", trunc(&msg_clone, 40));
}
#[cfg(not(windows))]
{
// On non-Windows just log it
println!(" [PROMPT] {}", message);
}
ClientMessage::Ack { request_id }
}
ServerMessage::ExecRequest { request_id, command, timeout_ms } => {
let cmd_display = trunc(&command, 60);
log_cmd!("", "exec {}", cmd_display);
let mut sh = shell.lock().await;
match sh.run(&command, timeout_ms).await {
Ok((stdout, stderr, exit_code)) => {
let out = stdout.trim().lines().next().unwrap_or("").to_string();
let out_display = trunc(&out, 60);
if exit_code != 0 {
if out_display.is_empty() {
log_err!("exit {}", exit_code);
} else {
log_err!("{} {} exit {}", out_display, "·".dimmed(), exit_code);
}
} else if out_display.is_empty() {
log_ok!("exit 0");
} else {
log_ok!("{} {} exit 0", out_display, "·".dimmed());
}
let _ = stderr;
ClientMessage::ExecResponse {
request_id,
stdout,
stderr,
exit_code,
}
}
Err(e) => {
log_err!("exec failed: {e}");
ClientMessage::Error {
request_id,
message: format!(
"Exec failed for command {:?}.\nError: {e}",
command
),
}
}
}
}
ServerMessage::ClickRequest { request_id, x, y, button } => {
log_cmd!("🖱 ", "click ({x}, {y}) {:?}", button);
match input::click(x, y, &button) {
Ok(()) => {
log_ok!("Done");
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("click failed: {e}");
ClientMessage::Error {
request_id,
message: format!("Click at ({x},{y}) failed: {e}"),
}
}
}
}
ServerMessage::TypeRequest { request_id, text } => {
log_cmd!("", "type {} chars", text.len());
match input::type_text(&text) {
Ok(()) => {
log_ok!("Done");
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("type failed: {e}");
ClientMessage::Error {
request_id,
message: format!("Type failed: {e}"),
}
}
}
}
ServerMessage::ListWindowsRequest { request_id } => {
log_cmd!("🪟", "list-windows");
match windows_mgmt::list_windows() {
Ok(windows) => {
log_ok!("{} windows", windows.len());
ClientMessage::ListWindowsResponse { request_id, windows }
}
Err(e) => {
log_err!("list-windows failed: {e}");
ClientMessage::Error { request_id, message: e }
}
}
}
ServerMessage::MinimizeAllRequest { request_id } => {
log_cmd!("🪟", "minimize-all");
match windows_mgmt::minimize_all() {
Ok(()) => {
log_ok!("Done");
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("minimize-all failed: {e}");
ClientMessage::Error { request_id, message: e }
}
}
}
ServerMessage::FocusWindowRequest { request_id, window_id } => {
log_cmd!("🪟", "focus-window {window_id}");
match windows_mgmt::focus_window(window_id) {
Ok(()) => {
log_ok!("Done");
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("focus-window failed: {e}");
ClientMessage::Error { request_id, message: e }
}
}
}
ServerMessage::MaximizeAndFocusRequest { request_id, window_id } => {
log_cmd!("🪟", "maximize-and-focus {window_id}");
match windows_mgmt::maximize_and_focus(window_id) {
Ok(()) => {
log_ok!("Done");
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("maximize-and-focus failed: {e}");
ClientMessage::Error { request_id, message: e }
}
}
}
ServerMessage::VersionRequest { request_id } => {
log_cmd!(" ", "version");
ClientMessage::VersionResponse {
request_id,
version: env!("CARGO_PKG_VERSION").to_string(),
commit: env!("GIT_COMMIT").to_string(),
}
}
ServerMessage::LogsRequest { request_id, lines } => {
log_cmd!("📜", "logs (last {lines} lines)");
let content = logger::tail(lines);
let log_path = logger::get_log_path();
ClientMessage::LogsResponse { request_id, content, log_path }
}
ServerMessage::UploadRequest { request_id, path, content_base64 } => {
log_cmd!("", "upload → {}", path);
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(()) => {
log_ok!("Saved {}", path);
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("upload failed: {e}");
ClientMessage::Error { request_id, message: e }
}
}
}
ServerMessage::DownloadRequest { request_id, path } => {
log_cmd!("", "download ← {}", path);
match std::fs::read(&path) {
Ok(bytes) => {
let size = bytes.len() as u64;
let content_base64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
log_ok!("Sent {} bytes", size);
ClientMessage::DownloadResponse { request_id, content_base64, size }
}
Err(e) => {
log_err!("download failed: {e}");
ClientMessage::Error { request_id, message: format!("Read failed: {e}") }
}
}
}
ServerMessage::RunRequest { request_id, program, args } => {
log_cmd!("🚀", "run {}", program);
use std::process::Command as StdCommand;
match StdCommand::new(&program).args(&args).spawn() {
Ok(_) => {
log_ok!("Started {}", program);
ClientMessage::Ack { request_id }
}
Err(e) => {
log_err!("run failed: {e}");
ClientMessage::Error { request_id, message: format!("Failed to start '{}': {e}", program) }
}
}
}
ServerMessage::ClipboardGetRequest { request_id } => {
log_cmd!("📋", "clipboard-get");
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();
log_ok!("Got {} chars", text.len());
ClientMessage::ClipboardGetResponse { request_id, text }
}
Err(e) => ClientMessage::Error { request_id, message: format!("Clipboard get failed: {e}") }
}
}
ServerMessage::ClipboardSetRequest { request_id, text } => {
log_cmd!("📋", "clipboard-set {} chars", text.len());
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(_) => {
log_ok!("Set clipboard");
ClientMessage::Ack { request_id }
}
Err(e) => ClientMessage::Error { request_id, message: format!("Clipboard set failed: {e}") }
}
}
ServerMessage::Ack { request_id } => {
ClientMessage::Ack { request_id }
}
ServerMessage::Error { request_id, message } => {
log_err!("server error: {message}");
if let Some(rid) = request_id {
ClientMessage::Ack { request_id: rid }
} else {
ClientMessage::Hello { label: None }
}
}
}
}