multi-ssh/src/main.rs

mod action;
mod file;
mod os_string_builder;
mod server;

use crate::action::{Action, FileAction, ServerActions};
use crate::file::{FileMatcher, FileNameInfo};
use crate::os_string_builder::ReplaceWithOsStr;
use clap::{Parser, Subcommand, ValueEnum};
use lazy_regex::{lazy_regex, Lazy, Regex};
use server::{Server, ServerReference};
use std::cell::LazyCell;
use std::hash::Hash;
use std::io::Write;
use std::iter::once;
use std::path::PathBuf;
use std::process::Stdio;
use std::str::FromStr;
use std::{env, fs};

const SERVERS_ENV_VAR: &str = "MSSH_SERVERS";
const EDITOR_ENV_VAR: &str = "MSSH_EDITOR";
const FILE_PLACEHOLDER: &str = "<file>";

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")]
    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')]
  #[allow(clippy::enum_variant_names)]
  Command {
    /// The command to execute
    command: String,
  },
  /// Open a file in your local editor and upload it afterward.
  ///
  /// The editor will be opened for every server in the order provided by you.
  /// The command to start the editor can be configured via environment variable `MSSH_EDITOR`:
  ///
  /// `export MSSH_EDITOR="nano <file>"`
  ///
  /// File is a placeholder that will be replaced with the actual name of the file.
  #[command(visible_short_flag_alias = 'e')]
  Editor {
    /// A path to the file to edit, relative to the server directory
    file: PathBuf,
    /// The command to start the editor. Supports the placeholder `<file>`, e.g. "nano <file>".
    ///
    /// If omitted, the command will be taken from the environment variable `MSSH_EDITOR`.
    #[arg(short, long)]
    editor: Option<String>,
    /// The directory where to save the file to
    #[arg(short = 'd', long, default_value = ".mssh/downloads")]
    working_directory: PathBuf,
    /// Override existing files in the working directory
    #[arg(short = 'f', long = "override", default_value = "false")]
    override_existing: bool,
  },
}

#[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,
    } => {
      if servers.is_empty() {
        println!("Please provide some servers to upload to. See --help");
        return Ok(());
      }

      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| {
          let working_directory = server.server_directory_path.join(&upload_directory);
          Ok(ServerActions {
            server,
            actions: {
              let output = ShellCmd::new("ssh")
                .arg(&server.ssh_name)
                .arg(osf!("ls ") + &working_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 => {
                  let mut actions: Vec<_> = files
                    .filter(|file| file_matcher.matches(file))
                    .map(|file| {
                      //special case -> file has the same name as current file, then we just need to replace it
                      if file == file_name {
                        FileAction {
                          file: PathBuf::from(file),
                          kind: Action::Replace,
                        }
                      } else {
                        FileAction {
                          file: PathBuf::from(file),
                          kind: Action::Delete,
                        }
                      }
                    })
                    .collect();
                  if !actions.iter().any(|action| action.kind == Action::Replace) {
                    actions.push(add_action);
                  }
                  actions
                }
              }
            },
            working_directory,
          })
        })
        .collect::<Result<Vec<_>, 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_actions.working_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_actions.working_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_actions.working_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!");
    }
    Command::Editor {
      file,
      editor,
      working_directory,
      override_existing,
    } => {
      //get editor
      let editor = editor.ok_or(()).or_else(|_| env::var(EDITOR_ENV_VAR).map_err(|e| format!("You have not specified an editor. Please do so using the --editor flag or the {EDITOR_ENV_VAR} environment variable: {e}")))?;

      fs::create_dir_all(&working_directory)
        .map_err(|e| format!("failed to create working directory: {e}"))?;

      //make sure file doesn't exist in working directory yet, or will be overridden
      let file_name = file
        .file_name()
        .ok_or("can only edit files, not directories")?;
      if !override_existing
        && fs::read_dir(&working_directory)
          .map_err(|e| format!("failed to open working directory: {e}"))?
          .collect::<Result<Vec<_>, _>>()
          .map_err(|e| format!("error while querying working directory contents: {e}"))?
          .iter()
          .any(|entry| entry.file_name() == file_name)
      {
        return Err(format!(
          "A file with the name {} already exists in {}. You can override it with --override",
          file_name.to_string_lossy(),
          working_directory.to_string_lossy()
        ));
      }

      start_ssh_agent()?;

      for server in servers {
        println!("Downloading file from {}...", server.ssh_name);
        ShellCmd::new("scp")
          .arg(osf!(&server.ssh_name) + ":" + server.server_directory_path.join(&file))
          .arg(&working_directory)
          .status()
          .map_err(|e| format!("failed to download file: {e}"))?;

        //open file in editor
        let mut shell_args = shell_words::split(&editor)
          .map_err(|e| format!("failed to parse editor command: {e}"))?
          .into_iter()
          .map(|part| part.replace_with_os_str(FILE_PLACEHOLDER, working_directory.join(file_name)))
          .collect::<Vec<_>>();

        let command = shell_args.remove(0);
        ShellCmd::new(command)
          .args(shell_args)
          .status()
          .map_err(|e| format!("failed to open file in editor: {e}"))?;

        //upload file again
        ShellCmd::new("scp")
          .arg(working_directory.join(file_name))
          .arg(osf!(&server.ssh_name) + ":" + server.server_directory_path.join(&file))
          .status()
          .map_err(|e| format!("failed to upload file again: {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_stdout = String::from_utf8_lossy(&agent_output.stdout);
  if !agent_output.status.success() {
    return Err("failed to start ssh agent; maybe try to run ssh-agent manually?".to_string());
  }

  //set the env vars from the agent
  static ENV_VAR_REGEX: Lazy<Regex> = lazy_regex!("(.+?)=(.+?);");
  for capture in ENV_VAR_REGEX.captures_iter(&agent_stdout) {
    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
    );
  }

  #[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);
  }
}