multi-ssh/src/main.rs

mod server;

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::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<ServerReference>,
}

#[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", default_missing_value = "true", num_args = 0..1)]
    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<String>,
  },
  /// 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::<Result<Vec<_>, _>>()?;

  match args.command {
    Command::Upload {
      file,
      old_version_policy,
      upload_directory,
      no_confirm,
      file_name,
    } => {
      todo!()
    }
    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(format!("cd {:?}; {command}", server.server_directory_path))
          .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<Regex> = 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<Vec<Server>, 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<Vec<Server>, 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
    );
  }
}