diff --git a/agent/src/service.rs b/agent/src/service.rs new file mode 100644 index 0000000..f59560c --- /dev/null +++ b/agent/src/service.rs @@ -0,0 +1,645 @@ +//! System service management for the Butterfly agent. +//! +//! Provides a cross-platform abstraction for installing, uninstalling, starting, +//! stopping, and querying the status of the agent as a background service. +//! +//! ## Platform Support +//! +//! - **Linux**: Uses `systemd` unit files. The service runs under the system +//! manager with configurable user and display environment variables. +//! - **Windows**: Uses the Windows Service Control Manager (SCM). The service +//! is registered via `sc.exe` commands and runs as a Windows Service with +//! proper lifecycle management (START/STOP/SHUTDOWN). +//! +//! ## Linux Systemd Unit File +//! +//! The generated unit file includes: +//! - Automatic restart on failure (with 5s delay) +//! - Display environment variable capture (DISPLAY, WAYLAND_DISPLAY, etc.) +//! - Configurable user and working directory +//! - `network-online.target` dependency for proper boot ordering + +use anyhow::{Context, Result}; +use log::{info, warn}; + +use crate::cli::ServiceAction; +use crate::config::RunOptions; + +/// Dispatch a service management action to the appropriate platform handler. +pub fn handle_service_action(action: ServiceAction) -> Result<()> { + match action { + ServiceAction::Install { + opts, + name, + display_name, + description, + start, + #[cfg(unix)] + user, + working_directory, + log_file, + } => install_service( + &name, + &display_name, + &description, + &opts, + #[cfg(unix)] + user.as_deref(), + working_directory.as_deref(), + log_file.as_deref(), + start, + ), + ServiceAction::Uninstall { name } => uninstall_service(&name), + ServiceAction::Start { name } => start_service(&name), + ServiceAction::Stop { name } => stop_service(&name), + ServiceAction::Status { name } => status_service(&name), + ServiceAction::Restart { name } => restart_service(&name), + } +} + +// ── Install ──────────────────────────────────────────────────────────────────── + +/// Install the agent as a system service. +/// +/// On Linux, this creates a systemd unit file and enables it. +/// On Windows, this registers with the Service Control Manager. +#[allow(clippy::too_many_arguments)] +fn install_service( + name: &str, + display_name: &str, + description: &str, + opts: &RunOptions, + #[cfg(unix)] user: Option<&str>, + working_directory: Option<&str>, + log_file: Option<&str>, + auto_start: bool, +) -> Result<()> { + #[cfg(unix)] + { + install_service_unix(name, display_name, description, opts, user, working_directory, log_file, auto_start) + } + #[cfg(windows)] + { + let _ = (user, working_directory, log_file); + install_service_windows(name, display_name, description, opts, auto_start) + } +} + +/// Install as a systemd service on Linux. +#[cfg(unix)] +#[allow(clippy::too_many_arguments)] +fn install_service_unix( + name: &str, + _display_name: &str, + description: &str, + opts: &RunOptions, + user: Option<&str>, + working_directory: Option<&str>, + log_file: Option<&str>, + auto_start: bool, +) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + + // Resolve the absolute path to the current binary. + let exe_path = std::env::current_exe() + .context("failed to resolve current executable path")?; + let exe_str = exe_path.to_str() + .context("executable path is not valid UTF-8")?; + + info!("installing systemd service '{}'", name); + info!(" binary: {}", exe_str); + info!(" server: {}", opts.server); + info!(" encoder: {}", opts.encoder); + info!(" fps: {}", opts.fps); + + // Build the ExecStart command line. + let args = opts.to_service_args(); + let exec_args: Vec = args.iter() + .flat_map(|a| { + // Quote arguments that contain spaces or special characters. + if a.contains(' ') || a.contains('"') || a.contains('\'') || a.contains('$') { + vec![format!("'{}'", a.replace('\'', "'\\''"))] + } else { + vec![a.clone()] + } + }) + .collect(); + let exec_start = format!("{} {}", exe_str, exec_args.join(" ")); + + // Detect display-related environment variables to capture. + let display_envs = detect_display_env(); + + // Build the systemd unit file. + let mut unit = String::new(); + + // [Unit] section. + unit.push_str(&format!("[Unit]\n")); + unit.push_str(&format!("Description={}\n", description)); + unit.push_str("After=network-online.target\n"); + unit.push_str("Wants=network-online.target\n"); + + // On systems with a graphical target, wait for it. + // This is optional — if the target doesn't exist, systemd ignores it. + unit.push_str("After=graphical.target\n"); + unit.push_str("\n"); + + // [Service] section. + unit.push_str("[Service]\n"); + unit.push_str("Type=simple\n"); + unit.push_str(&format!("ExecStart={}\n", exec_start)); + unit.push_str("Restart=always\n"); + unit.push_str("RestartSec=5\n"); + + // User. + if let Some(u) = user { + unit.push_str(&format!("User={}\n", u)); + unit.push_str(&format!("Group={}\n", u)); + } + + // Working directory. + if let Some(wd) = working_directory { + unit.push_str(&format!("WorkingDirectory={}\n", wd)); + } + + // Environment variables. + unit.push_str("Environment=RUST_LOG=info\n"); + for (key, value) in &display_envs { + unit.push_str(&format!("Environment={}={}\n", key, value)); + } + + // Logging to file. + if let Some(log_path) = log_file { + unit.push_str(&format!("StandardOutput=append:{}", log_path)); + unit.push_str(&format!("StandardError=append:{}", log_path)); + } + + // Security: don't kill the service when the user logs out. + // Important for headless VMs accessed via RDP/SSH. + unit.push_str("\n# Prevent service from being killed on user logout\n"); + unit.push_str("KillMode=process\n"); + + unit.push_str("\n[Install]\n"); + unit.push_str("WantedBy=multi-user.target\n"); + + // Write the unit file. + let unit_path = format!("/etc/systemd/system/{}.service", name); + info!("writing unit file: {}", unit_path); + + std::fs::write(&unit_path, &unit) + .context(format!("failed to write unit file to '{}'. Do you need sudo?", unit_path))?; + + // Set permissions (644). + let perms = std::fs::Permissions::from_mode(0o644); + std::fs::set_permissions(&unit_path, perms)?; + + // Reload systemd daemon. + info!("reloading systemd daemon..."); + let status = std::process::Command::new("systemctl") + .arg("daemon-reload") + .status() + .context("failed to run systemctl daemon-reload")?; + + if !status.success() { + warn!("systemctl daemon-reload returned non-zero exit code"); + } + + // Enable and optionally start the service. + if auto_start { + info!("enabling and starting service '{}'...", name); + let status = std::process::Command::new("systemctl") + .args(["enable", "--now", name]) + .status() + .context(format!("failed to enable/start service '{}'", name))?; + + if !status.success() { + anyhow::bail!( + "systemctl enable --now {} failed. Check 'systemctl status {}' for details.", + name, name + ); + } + } else { + info!("enabling service '{}' (not starting)...", name); + let status = std::process::Command::new("systemctl") + .args(["enable", name]) + .status() + .context(format!("failed to enable service '{}'", name))?; + + if !status.success() { + anyhow::bail!("systemctl enable {} failed", name); + } + } + + println!(); + println!("✅ Service '{}' installed successfully!", name); + println!(" Unit file: {}", unit_path); + println!(); + println!(" Commands:"); + println!(" sudo systemctl start {} — start the service", name); + println!(" sudo systemctl stop {} — stop the service", name); + println!(" sudo systemctl restart {} — restart the service", name); + println!(" sudo systemctl status {} — check status", name); + println!(" sudo journalctl -u {} -f — follow logs", name); + println!(); + + Ok(()) +} + +/// Install as a Windows Service. +#[cfg(windows)] +fn install_service_windows( + name: &str, + display_name: &str, + description: &str, + opts: &RunOptions, + auto_start: bool, +) -> Result<()> { + use std::process::Command; + + // Resolve the absolute path to the current binary. + let exe_path = std::env::current_exe() + .context("failed to resolve current executable path")?; + let exe_str = exe_path.to_str() + .context("executable path is not valid UTF-8")?; + + info!("installing Windows service '{}'", name); + info!(" binary: {}", exe_str); + info!(" server: {}", opts.server); + info!(" encoder: {}", opts.encoder); + + // Build the full command line for the service binary. + let args = opts.to_service_args(); + // On Windows, add the --windows-service flag so the binary knows to + // connect to the SCM dispatcher instead of running in foreground mode. + let mut all_args: Vec = args; + all_args.push("--windows-service".to_string()); + + // Quote the binary path (important if it contains spaces). + let bin_path = format!("\"{}\" {}", exe_str, all_args.join(" ")); + + // Use sc.exe to create the service. + // The sc.exe command has unusual argument parsing: key=value pairs where + // the value extends until the next recognized key. The space after '=' is + // significant and required. + let sc_args = [ + "create", + name, + &format!("binPath= \"{}\"", bin_path), + &format!("DisplayName= \"{}\"", display_name), + "start= auto", + "obj= LocalSystem", + "type= own", + ]; + + info!("running sc create {}...", name); + let result = Command::new("sc") + .args(&sc_args) + .output() + .context("failed to execute sc.exe. Are you running as Administrator?")?; + + if !result.status.success() { + let stdout = String::from_utf8_lossy(&result.stdout); + let stderr = String::from_utf8_lossy(&result.stderr); + anyhow::bail!( + "sc create failed:\n stdout: {}\n stderr: {}", + stdout.trim(), + stderr.trim() + ); + } + + // Set the service description. + let desc_result = Command::new("sc") + .args(["description", name, description]) + .output(); + + match desc_result { + Ok(r) if r.status.success() => {} + Ok(r) => warn!("failed to set service description: {}", String::from_utf8_lossy(&r.stderr).trim()), + Err(e) => warn!("failed to run sc description: {}", e), + } + + // Set failure recovery options: restart after 5 seconds, reset after 1 day. + let failure_result = Command::new("sc") + .args([ + "failure", + name, + "reset= 86400", + "actions= restart/5000", + ]) + .output(); + + match failure_result { + Ok(r) if r.status.success() => info!("configured failure recovery (restart after 5s)"), + Ok(r) => warn!("failed to set failure actions: {}", String::from_utf8_lossy(&r.stderr).trim()), + Err(e) => warn!("failed to set failure actions: {}", e), + } + + // Optionally start the service. + if auto_start { + info!("starting service '{}'...", name); + let start_result = Command::new("sc") + .args(["start", name]) + .output() + .context("failed to start service")?; + + if !start_result.status.success() { + let stdout = String::from_utf8_lossy(&start_result.stdout); + let stderr = String::from_utf8_lossy(&start_result.stderr); + warn!( + "service may not have started:\n stdout: {}\n stderr: {}", + stdout.trim(), + stderr.trim() + ); + } + } + + println!(); + println!("✅ Service '{}' installed successfully!", name); + println!(); + println!(" Commands:"); + println!(" sc start {} — start the service", name); + println!(" sc stop {} — stop the service", name); + println!(" sc query {} — check status", name); + println!(" sc delete {} — uninstall the service", name); + println!(); + println!(" Note: The service runs as LocalSystem. If screen capture"); + println!(" doesn't work, you may need to configure it to run as your"); + println!(" user account: sc config {} obj= \".\\USERNAME\" pwd= PASSWORD", name); + println!(); + + Ok(()) +} + +// ── Uninstall ────────────────────────────────────────────────────────────────── + +/// Uninstall the system service. +fn uninstall_service(name: &str) -> Result<()> { + #[cfg(unix)] + { + uninstall_service_unix(name) + } + #[cfg(windows)] + { + uninstall_service_windows(name) + } +} + +#[cfg(unix)] +fn uninstall_service_unix(name: &str) -> Result<()> { + let unit_path = format!("/etc/systemd/system/{}.service", name); + + // Check if the unit file exists. + if !std::path::Path::new(&unit_path).exists() { + anyhow::bail!("service '{}' not found (no unit file at {})", name, unit_path); + } + + info!("stopping service '{}'...", name); + let _ = std::process::Command::new("systemctl") + .args(["stop", name]) + .status(); + + info!("disabling service '{}'...", name); + let _ = std::process::Command::new("systemctl") + .args(["disable", name]) + .status(); + + info!("removing unit file '{}'...", unit_path); + std::fs::remove_file(&unit_path) + .context(format!("failed to remove unit file '{}'. Do you need sudo?", unit_path))?; + + info!("reloading systemd daemon..."); + let _ = std::process::Command::new("systemctl") + .arg("daemon-reload") + .status(); + + // Reset failed state if any. + let _ = std::process::Command::new("systemctl") + .args(["reset-failed", name]) + .status(); + + println!("✅ Service '{}' uninstalled successfully!", name); + Ok(()) +} + +#[cfg(windows)] +fn uninstall_service_windows(name: &str) -> Result<()> { + use std::process::Command; + + // Stop the service first. + info!("stopping service '{}'...", name); + let _ = Command::new("sc").args(["stop", name]).output(); + + // Delete the service. + info!("deleting service '{}'...", name); + let result = Command::new("sc") + .args(["delete", name]) + .output() + .context("failed to execute sc delete. Are you running as Administrator?")?; + + if !result.status.success() { + let stderr = String::from_utf8_lossy(&result.stderr); + anyhow::bail!("sc delete failed: {}", stderr.trim()); + } + + println!("✅ Service '{}' uninstalled successfully!", name); + Ok(()) +} + +// ── Start ────────────────────────────────────────────────────────────────────── + +/// Start the installed service. +fn start_service(name: &str) -> Result<()> { + #[cfg(unix)] + { + let result = std::process::Command::new("systemctl") + .args(["start", name]) + .status() + .context(format!("failed to start service '{}'. Do you need sudo?", name))?; + + if result.success() { + println!("✅ Service '{}' started.", name); + } else { + anyhow::bail!("failed to start service '{}'", name); + } + } + #[cfg(windows)] + { + let result = std::process::Command::new("sc") + .args(["start", name]) + .output() + .context("failed to start service. Are you running as Administrator?")?; + + if result.status.success() { + println!("✅ Service '{}' started.", name); + } else { + let stderr = String::from_utf8_lossy(&result.stderr); + anyhow::bail!("failed to start service '{}': {}", name, stderr.trim()); + } + } + Ok(()) +} + +// ── Stop ─────────────────────────────────────────────────────────────────────── + +/// Stop the running service. +fn stop_service(name: &str) -> Result<()> { + #[cfg(unix)] + { + let result = std::process::Command::new("systemctl") + .args(["stop", name]) + .status() + .context(format!("failed to stop service '{}'. Do you need sudo?", name))?; + + if result.success() { + println!("✅ Service '{}' stopped.", name); + } else { + anyhow::bail!("failed to stop service '{}'", name); + } + } + #[cfg(windows)] + { + let result = std::process::Command::new("sc") + .args(["stop", name]) + .output() + .context("failed to stop service. Are you running as Administrator?")?; + + if result.status.success() { + println!("✅ Service '{}' stopped.", name); + } else { + let stderr = String::from_utf8_lossy(&result.stderr); + anyhow::bail!("failed to stop service '{}': {}", name, stderr.trim()); + } + } + Ok(()) +} + +// ── Restart ──────────────────────────────────────────────────────────────────── + +/// Restart the running service. +fn restart_service(name: &str) -> Result<()> { + #[cfg(unix)] + { + let result = std::process::Command::new("systemctl") + .args(["restart", name]) + .status() + .context(format!("failed to restart service '{}'. Do you need sudo?", name))?; + + if result.success() { + println!("✅ Service '{}' restarted.", name); + } else { + anyhow::bail!("failed to restart service '{}'", name); + } + } + #[cfg(windows)] + { + let _ = stop_service(name); + // Brief pause to let the service fully stop. + std::thread::sleep(std::time::Duration::from_secs(2)); + start_service(name)?; + } + Ok(()) +} + +// ── Status ───────────────────────────────────────────────────────────────────── + +/// Check and display the service status. +fn status_service(name: &str) -> Result<()> { + #[cfg(unix)] + { + let result = std::process::Command::new("systemctl") + .args(["status", name]) + .status(); + + match result { + Ok(status) => { + // systemctl status returns 0 for active, 3 for inactive, etc. + // Let's also run is-active for a cleaner check. + let active = std::process::Command::new("systemctl") + .args(["is-active", name]) + .output(); + + match active { + Ok(output) => { + let state = String::from_utf8_lossy(&output.stdout).trim().to_string(); + if state == "active" { + println!("✅ Service '{}' is running.", name); + } else if state == "inactive" { + println!("⏹ Service '{}' is stopped.", name); + } else if state == "failed" { + println!("❌ Service '{}' has failed.", name); + println!(" Run 'systemctl status {}' for details.", name); + } else { + println!("❓ Service '{}' status: {}", name, state); + } + } + Err(_) => { + println!("❓ Could not determine service status (exit code: {:?})", status.code()); + } + } + } + Err(e) => { + anyhow::bail!("failed to query service '{}': {}", name, e); + } + } + } + #[cfg(windows)] + { + let result = std::process::Command::new("sc") + .args(["query", name]) + .output() + .context("failed to query service")?; + + if result.status.success() { + let stdout = String::from_utf8_lossy(&result.stdout); + println!("{}", stdout.trim()); + } else { + let stderr = String::from_utf8_lossy(&result.stderr); + anyhow::bail!("failed to query service '{}': {}", name, stderr.trim()); + } + } + Ok(()) +} + +// ── Helpers ──────────────────────────────────────────────────────────────────── + +/// Detect display-related environment variables for embedding in the systemd unit. +/// +/// These variables are needed for screen capture and input injection to work +/// from within a systemd service (which otherwise runs with a minimal environment). +#[cfg(unix)] +fn detect_display_env() -> Vec<(String, String)> { + let mut envs = Vec::new(); + + // X11 display server. + if let Ok(val) = std::env::var("DISPLAY") { + envs.push(("DISPLAY".into(), val)); + } + + // Xauthority file for X11 authentication. + if let Ok(val) = std::env::var("XAUTHORITY") { + envs.push(("XAUTHORITY".into(), val)); + } + + // Wayland display server. + if let Ok(val) = std::env::var("WAYLAND_DISPLAY") { + envs.push(("WAYLAND_DISPLAY".into(), val)); + } + + // Wayland runtime directory (needed for Wayland socket access). + if let Ok(val) = std::env::var("XDG_RUNTIME_DIR") { + envs.push(("XDG_RUNTIME_DIR".into(), val)); + } + + // D-Bus session bus (needed for some desktop interactions). + if let Ok(val) = std::env::var("DBUS_SESSION_BUS_ADDRESS") { + envs.push(("DBUS_SESSION_BUS_ADDRESS".into(), val)); + } + + if envs.is_empty() { + log::warn!("no display environment variables detected — the service may not be able to capture the screen"); + log::warn!("run this from a desktop session, or manually set DISPLAY=:0 in the unit file"); + } else { + info!("captured display environment: {:?}", envs.iter().map(|(k, _)| k.as_str()).collect::>()); + } + + envs +}