mod action; mod file; mod os_string_builder; mod server; use crate::action::{Action, FileAction, ServerActions}; use crate::file::{FileMatcher, FileNameInfo}; use clap::{Parser, Subcommand, ValueEnum}; use lazy_regex::{lazy_regex, Lazy, Regex}; use server::{Server, ServerReference}; use std::cell::LazyCell; use std::env; use std::hash::Hash; use std::io::Write; use std::iter::once; use std::path::PathBuf; use std::process::Stdio; use std::str::FromStr; const SERVERS_ENV_VAR: &str = "MSSH_SERVERS"; type ShellCmd = std::process::Command; /// Uploads a file or executes a command on multiple configured servers /// /// Servers must either be configured via environment variable or denote their server directory with /// a double colon: crea:home/crea. /// /// --- Configuration via environment variable --- /// /// Use MSSH_SERVERS="crea:home/crea,sky:sky,lobby:,city:city2" to configure servers. #[derive(Parser, Debug)] #[command(version, about, long_about)] struct Args { /// The action to perform #[command(subcommand)] command: Command, /// The ssh names and optionally home directories of the servers to perform the action on #[arg(num_args = 1.., value_parser = ServerReference::from_str)] servers: Vec, } #[derive(Subcommand, Debug)] enum Command { /// Upload a file to the servers #[command(visible_short_flag_alias = 'u')] Upload { /// The file to upload file: PathBuf, /// How to handle older versions of the file #[arg(short = 'a', long, default_value = "delete", default_missing_value = "archive", num_args = 0..=1)] old_version_policy: OldVersionPolicy, /// The directory where to upload to, relative to the server directory #[arg(short = 'p', long, default_value = "plugins")] upload_directory: PathBuf, /// Skip the confirmation dialog #[arg(long, default_value = "false")] no_confirm: bool, /// The prefix of the name of older versions of the file, which should be replaced or deleted #[arg(short, long)] file_name: Option, }, /// Execute a command on the servers #[command(visible_short_flag_alias = 'c')] Command { /// The command to execute command: String, }, } #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash, Default, ValueEnum)] enum OldVersionPolicy { /// Ignore the existence of older versions Ignore, /// Rename older versions: foo.jar -> foo.jarr Archive, /// Delete older versions #[default] Delete, } fn main() -> Result<(), String> { let args = Args::parse(); let mut configured_servers = LazyCell::new(parse_server_configuration_from_env); let servers = args .servers .iter() .map(|server_reference| { let server_name = server_reference.get_name(); server_reference .clone() .try_resolve_lazy(&mut configured_servers) .map_err(|msg| format!("Can't resolve server directory for '{server_name}': {msg}")) .and_then(|opt_server| { opt_server.ok_or(format!( "no server directory has been configured for server '{server_name}'" )) }) }) .collect::, _>>()?; match args.command { Command::Upload { file, old_version_policy, upload_directory, no_confirm, file_name, } => { start_ssh_agent()?; let file_name_info = FileNameInfo::try_from(file.clone()).map_err(|e| format!("bad file: {e}"))?; //create overview of what has to be done on each server let actions = servers .iter() .map(|server| { Ok(ServerActions { server, actions: { let output = ShellCmd::new("ssh") .arg(&server.ssh_name) .arg(osf!("ls ") + server.server_directory_path.join(&upload_directory)) .stdout(Stdio::piped()) .output() .map_err(|e| format!("failed to query files via ssh: {e}"))?; let output = String::from_utf8_lossy(&output.stdout); let mut file_matcher = FileMatcher::from(file_name.as_ref().unwrap_or(&file_name_info.name)); if let Some(extension) = file_name_info.extension.as_ref() { file_matcher = file_matcher.and_extension(extension); } let file_name = file_name_info.to_full_file_name(); let add_action = FileAction { file: PathBuf::from(&file_name), kind: Action::Add, }; let mut files = output.lines(); match old_version_policy { OldVersionPolicy::Ignore => { vec![if files.any(|file| file == file_name) { FileAction { file: PathBuf::from(&file_name), kind: Action::Replace, } } else { add_action }] } OldVersionPolicy::Archive => files .filter(|file| file_matcher.matches(file)) .map(|file| FileAction { file: PathBuf::from(file), kind: Action::Rename { new_name: format!("{file}{}", file.chars().last().unwrap_or('1')).into(), }, }) .chain(once(add_action)) .collect(), OldVersionPolicy::Delete => files .filter(|file| file_matcher.matches(file)) .map(|file| FileAction { file: PathBuf::from(file), kind: Action::Delete, }) .chain(once(add_action)) .collect(), } }, }) }) .collect::, String>>()?; println!("The following actions will be performed:"); for server_actions in &actions { println!("{server_actions}"); } if !no_confirm { print!("Continue? [Y|n] "); std::io::stdout().flush().expect("failed to flush stdout"); let mut buffer = String::new(); std::io::stdin() .read_line(&mut buffer) .expect("failed to read stdin"); match buffer.to_lowercase().trim() { "n" | "no" => { println!("Aborting..."); return Ok(()); } _ => {} } } for server_actions in actions { let server = server_actions.server; println!("Performing actions on {}...", server.ssh_name); for file_action in server_actions.actions { match file_action.kind { Action::Add | Action::Replace => { ShellCmd::new("scp") .arg(file.clone()) .arg( osf!(&server.ssh_name) + ":" + server.server_directory_path.join(&upload_directory), ) .spawn() .map_err(|e| format!("failed to upload file: {e}"))? .wait() .map_err(|e| format!("failed to wait for upload: {e}"))?; } Action::Delete => { ShellCmd::new("ssh") .arg(&server.ssh_name) .arg( osf!("cd ") + server.server_directory_path.join(&upload_directory) + "; rm " + &file_action.file, ) .spawn() .map_err(|e| format!("failed to send delete command: {e}"))? .wait() .map_err(|e| format!("failed to wait for delete command: {e}"))?; } Action::Rename { new_name } => { ShellCmd::new("ssh") .arg(&server.ssh_name) .arg( osf!("cd ") + server.server_directory_path.join(&upload_directory) + "; mv " + &file_action.file + " " + new_name, ) .spawn() .map_err(|e| format!("failed to send rename command: {e}"))? .wait() .map_err(|e| format!("failed to wait for rename command: {e}"))?; } } } } println!("Done!"); } Command::Command { command } => { start_ssh_agent()?; for server in servers { println!("Running command on '{}'...", server.ssh_name); ShellCmd::new("ssh") .arg(server.ssh_name) .arg(osf!("cd ") + server.server_directory_path + "; " + &command) .spawn() .map_err(|_| "failed to start ssh command".to_string())? .wait() .map_err(|e| format!("failed to wait for ssh command completion: {e}"))?; } println!("Done!"); } } Ok(()) } fn start_ssh_agent() -> Result<(), String> { //start the ssh agent let agent_output = ShellCmd::new("ssh-agent") .arg("-s") .stdout(Stdio::piped()) .output() .map_err(|e| format!("failed to start ssh agent: {e}"))?; let agent_output = String::from_utf8_lossy(&agent_output.stdout); //set the env vars from the agent static ENV_VAR_REGEX: Lazy = lazy_regex!("(.+?)=(.+?);"); for capture in ENV_VAR_REGEX.captures_iter(&agent_output) { let (_, [env_var, value]) = capture.extract(); env::set_var(env_var, value); } //add the ssh key ShellCmd::new("ssh-add") .spawn() .map_err(|e| format!("failed to add ssh key: {}", e))? .wait() .expect("failed to wait on ssh-add"); Ok(()) } fn parse_server_configuration_from_env() -> Result, String> { env::var(SERVERS_ENV_VAR) .map_err(|_| format!("Missing environment variable {}", SERVERS_ENV_VAR)) .and_then(|value| parse_server_configuration(&value)) } fn parse_server_configuration(config_str: &str) -> Result, String> { config_str .split(',') .map(|server_entry| { Server::from_str(server_entry) .map_err(|e| format!("Invalid server entry '{server_entry}': {e}")) }) .collect() } #[cfg(test)] mod test { use crate::parse_server_configuration; use crate::server::Server; use std::path::PathBuf; #[test] fn test_parse_server_configuration() { let servers = parse_server_configuration("foo:bar,fizz:buzz/bizz").expect("valid server configuration"); assert_eq!( vec![ Server { ssh_name: "foo".to_string(), server_directory_path: PathBuf::from("bar"), }, Server { ssh_name: "fizz".to_string(), server_directory_path: PathBuf::from("buzz/bizz"), } ], servers ); } #[test] fn path_experiment() { let server_dir = PathBuf::from("steptech"); let upload_dir = PathBuf::from("/home"); //absolute path let joined = server_dir.join(upload_dir); assert_eq!(PathBuf::from("/home"), joined); } }