multi-ssh/src/main.rs

mod action;
mod command;
mod file;
mod logger;
mod os_string_builder;
mod server;

use crate::action::{Action, FileAction, ServerActions};
use crate::command::{ExecutionError, LogRunnable, SpecificExecutionError};
use crate::file::{FileMatcher, FileNameInfo};
use crate::logger::{LogLevel, Logger};
use crate::os_string_builder::ReplaceWithOsStr;
use crate::server::{RelativeLocalPathAnker, ServerAddress};
use clap::{Parser, Subcommand, ValueEnum};
use lazy_regex::{lazy_regex, Lazy, Regex};
use server::{Server, ServerReference};
use std::cell::LazyCell;
use std::ffi::{OsStr, OsString};
use std::hash::Hash;
use std::io::Write;
use std::iter::once;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::{env, fs, io};

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>,
  /// How verbose logging output should be
  #[arg(long, default_value = "info", conflicts_with_all = ["quiet", "info"])]
  log_level: LogLevel,
  /// Only log errors
  #[arg(short, long, default_value = "false", conflicts_with_all = ["info"])]
  quiet: bool,
  /// Log additional debugging info
  #[arg(short = 'v', long, default_value = "false")]
  info: bool,
}

#[derive(Subcommand, Debug)]
enum Command {
  /// Upload a file to the servers
  #[command(visible_short_flag_alias = 'u')]
  Upload {
    /// The files to upload
    files: Vec<PathBuf>,
    /// The ssh server to get the file from.
    ///
    /// When this option is set, the file path must be absolute, or relative to the server directory.
    /// The upload-directory has no influence on where the file will be taken from.
    #[arg(short = 'S', long)]
    file_server: Option<ServerReference>,
    /// 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.
    ///
    /// Default directory is `~/Downloads`
    #[arg(short = 'd', long)]
    download_directory: Option<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,
}

#[macro_export]
macro_rules! input {
  ($prompt: tt) => {{
    print!($prompt);
    io::stdout().flush().expect("failed to flush stdout");
    let mut buf = String::new();
    io::stdin()
      .read_line(&mut buf)
      .expect("failed to read stdin");
    buf.trim().to_string()
  }};
  () => {
    input!()
  };
}

fn main() -> Result<(), String> {
  let args = Args::parse();

  let logger = Logger {
    //all the below options are conflicting with each other so an if else is fine
    level: if args.quiet {
      LogLevel::Error
    } else if args.info {
      LogLevel::Debug
    } else {
      args.log_level
    },
  };

  let mut configured_servers = LazyCell::new(parse_server_configuration_from_env);
  let servers = args
    .servers
    .iter()
    .map(|server_reference| {
      let server_identifier = server_reference.get_identifier();
      server_reference
        .clone()
        .try_resolve_lazy(&mut configured_servers)
        .map_err(|msg| format!("Can't resolve server directory for '{server_identifier}': {msg}"))
        .and_then(|opt_server| {
          opt_server.ok_or(format!(
            "no server directory has been configured for server '{server_identifier}'"
          ))
        })
    })
    .collect::<Result<Vec<_>, _>>()?;

  match args.command {
    Command::Upload {
      mut files,
      file_server,
      old_version_policy,
      mut upload_directory,
      no_confirm,
      file_name,
    } => {
      require_non_empty_servers(&servers)?;
      require_non_empty(&files, "files to upload")?;
      start_ssh_agent(&logger)?;

      //resolve file server
      let file_server = match file_server {
        Some(server_reference) => {
          let file_server_identifier = server_reference.get_identifier().to_string();
          let server = server_reference.try_resolve_lazy(&mut configured_servers)
            .map_err(|e| format!("Can't resolve server directory for file-server '{file_server_identifier}': {e}"))?
            .ok_or_else(|| format!("no server directory has been configured for file-server '{file_server_identifier}'"))?;
          Some(server)
        }
        None => None,
      };

      //make sure files exist
      match &file_server {
        Some(file_server) => match &file_server.address {
          ServerAddress::Ssh { ssh_address } => {
            //canonicalize remote files
            //TODO only join when there are multiple paths
            let joined_files = osf!("{{")
              + files
                .iter()
                .map(|file| file.as_os_str())
                .collect::<Vec<_>>()
                .join(&OsString::from(","))
              + "}";
            
            //TODO handle bad exit status
            let realpath_output = ShellCmd::new("ssh")
              .arg(ssh_address)
              .arg(osf!("realpath ") + file_server.server_directory_path.join(joined_files))
              .collect_output()
              .map_err(|e| format!("failed to canonicalize files to upload: {e}"))?;
            files = realpath_output
              .stdout
              .split(|&b| b == b'\n') //split at line breaks
              .map(|bytes| PathBuf::from(OsStr::from_bytes(bytes)))
              .collect();
            log!(logger, debug, "canonical files: {files:?}");

            for file in &files {
              check_file_exists_on_server(file, ssh_address, &file_server.server_directory_path)?;
            }
          }
          ServerAddress::Localhost => {
            for file in &files {
              check_local_file_exists(file_server.server_directory_path.join(file))?;
            }
          }
        },
        None => {
          for file in &files {
            check_local_file_exists(file)?;
          }
        }
      }

      let file_details = files
        .iter()
        .map(|file| {
          FileNameInfo::try_from(file.clone())
            .map(|info| (PathBuf::from(file), info))
            .map_err(|e| format!("Bad file '{}': {e}", file.to_string_lossy()))
        })
        .collect::<Result<Vec<_>, _>>()?;

      //create overview of what has to be done on each server
      let actions = servers
        .iter()
        .map(|server| {
          //on local server canonicalize upload_directory
          if let ServerAddress::Localhost = &server.address {
            //create upload directory if it doesn't exist
            fs::create_dir_all(&upload_directory)
              .map_err(|e| format!("Failed to create upload-directory: {e}"))?;

            upload_directory = fs::canonicalize(&upload_directory)
              .map_err(|e| format!("failed to resolve upload-directory: {e}"))?;
          }
          let working_directory = server.server_directory_path.join(&upload_directory);
          Ok(ServerActions {
            server,
            actions: {
              let mut ls_command = match &server.address {
                ServerAddress::Ssh { ssh_address } => {
                  let mut cmd = ShellCmd::new("ssh");
                  cmd.arg(ssh_address).arg(osf!("ls ") + &working_directory);
                  cmd
                }
                ServerAddress::Localhost => {
                  let mut cmd = ShellCmd::new("ls");
                  cmd.arg(&working_directory);
                  cmd
                }
              };
              let ls_output = ls_command
                .collect_output()
                .map_err(|e| format!("failed to query files: {e}"))?;
              let ls_output = String::from_utf8_lossy(&ls_output.stdout);

              file_details
                .iter()
                .flat_map(|(file, file_name_info)| {
                  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::new(file, Action::Add).expect("path points to file");

                  let mut ls_lines = ls_output.lines();
                  match old_version_policy {
                    OldVersionPolicy::Ignore => {
                      vec![if ls_lines.any(|file| file == file_name) {
                        FileAction::new(&file_name, Action::Replace).expect("path points to file")
                      } else {
                        add_action
                      }]
                    }
                    OldVersionPolicy::Archive => ls_lines
                      .filter(|file| file_matcher.matches(file))
                      .map(|file| {
                        FileAction::new(
                          file,
                          Action::rename(format!("{file}{}", file.chars().last().unwrap_or('1'))),
                        )
                        .expect("path points to file")
                      })
                      .chain(once(add_action))
                      .collect(),
                    OldVersionPolicy::Delete => {
                      let mut actions = ls_lines
                        .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::new(file, Action::Replace).expect("path points to file")
                          } else {
                            FileAction::new(file, Action::Delete).expect("path points to file")
                          }
                        })
                        .collect::<Vec<_>>();
                      if !actions.iter().any(|action| action.kind == Action::Replace) {
                        actions.push(add_action);
                      }
                      actions
                    }
                  }
                })
                .collect()
            },
            working_directory,
          })
        })
        .collect::<Result<Vec<_>, String>>()?;

      log!(logger, "The following actions will be performed: ");
      for server_actions in &actions {
        log!(logger, "{server_actions}");
      }

      if !no_confirm {
        match input!("Continue? [Y|n] ").to_lowercase().as_str() {
          "n" | "no" => {
            log!(logger, "Aborting...");
            return Ok(());
          }
          _ => {}
        }
      }

      for server_actions in actions {
        let server = server_actions.server;
        log!(logger, "Performing actions on {}...", server.get_name());
        for file_action in server_actions.actions {
          match file_action.kind {
            Action::Add | Action::Replace => {
              let scp_source = match &file_server {
                Some(file_server) => {
                  osf!(match &file_server.address {
                    ServerAddress::Ssh { ssh_address } => format!("{ssh_address}:"),
                    ServerAddress::Localhost => "".to_string(),
                  }) + file_server.server_directory_path.join(&file_action.file)
                }
                None => osf!(&file_action.file),
              };
              let scp_target = osf!(match &server.address {
                ServerAddress::Ssh { ssh_address } => format!("{ssh_address}:"),
                ServerAddress::Localhost => "".to_string(),
              }) + &server_actions.working_directory;
              ShellCmd::new("scp")
                .arg(scp_source)
                .arg(scp_target)
                .run(&logger)
                .map_err(|e| format!("upload failure: {e}"))?;
            }
            Action::Delete => match &server.address {
              ServerAddress::Ssh { ssh_address } => {
                ShellCmd::new("ssh")
                  .arg(ssh_address)
                  .arg(osf!("rm ") + server_actions.working_directory.join(&file_action.file))
                  .run(&logger)
                  .map_err(|e| format!("failed to delete old version: {e}"))?;
              }
              ServerAddress::Localhost => {
                fs::remove_file(server_actions.working_directory.join(&file_action.file))
                  .map_err(|e| format!("failed to delete old version: {e}"))?;
              }
            },
            Action::Rename { new_name } => match &server.address {
              ServerAddress::Ssh { ssh_address } => {
                ShellCmd::new("ssh")
                  .arg(ssh_address)
                  .arg(
                    osf!("mv ")
                      + server_actions.working_directory.join(&file_action.file)
                      + " "
                      + server_actions.working_directory.join(&new_name),
                  )
                  .run(&logger)
                  .map_err(|e| format!("failed to rename: {e}"))?;
              }
              ServerAddress::Localhost => {
                let dir = &server_actions.working_directory;
                fs::rename(dir.join(&file_action.file), dir.join(&new_name))
                  .map_err(|e| format!("failed to rename: {e}"))?;
              }
            },
          }
        }
      }

      log!(logger, "Done!");
    }
    Command::Command { command } => {
      start_ssh_agent(&logger)?;
      require_non_empty_servers(&servers)?;
      for server in servers {
        log!(logger, "Running command on '{}'...", server.get_name());
        match &server.address {
          ServerAddress::Ssh { ssh_address } => {
            ShellCmd::new("ssh")
              .arg(ssh_address)
              .arg(osf!("cd ") + server.server_directory_path + "; " + &command)
              .run(&logger)
              .map_err(|e| format!("{e}"))?;
          }
          ServerAddress::Localhost => {
            let mut command_args =
              shell_words::split(&command).map_err(|e| format!("failed to parse command: {e}"))?;
            ShellCmd::new(command_args.remove(0))
              .args(&command_args)
              .current_dir(&server.server_directory_path)
              .run(&logger)
              .map_err(|e| format!("{e}"))?;
          }
        }
      }
      log!(logger, "Done!");
    }
    Command::Editor {
      file,
      editor,
      download_directory,
      override_existing,
    } => {
      //determine download directory
      let download_directory = match download_directory {
        Some(download_directory) => download_directory,
        None => {
          let home_dir =
            get_home_directory().map_err(|e| format!("Can't determine download directory: {e}"))?;
          home_dir.join("Downloads")
        }
      };

      //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(&download_directory)
        .map_err(|e| format!("failed to create working directory: {e}"))?;

      //make sure file doesn't exist in working directory yet, or it will be overridden
      let file_name = file
        .file_name()
        .ok_or("can only edit files, not directories")?;
      'duplicate_check: {
        if !override_existing
          && fs::read_dir(&download_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)
        {
          let duplication_notification = format!(
            "A file with the name {} already exists in {}",
            file_name.to_string_lossy(),
            download_directory.to_string_lossy()
          );

          if !args.quiet {
            match input!("{duplication_notification}. Do you want to replace it? [N|y] ")
              .to_lowercase()
              .as_str()
            {
              "y" | "yes" => break 'duplicate_check,
              _ => {}
            }
          }

          return Err(format!(
            "{duplication_notification}. You can override it with --override or -f"
          ));
        }
      }

      require_non_empty_servers(&servers)?;
      start_ssh_agent(&logger)?;

      for server in servers {
        log!(logger, "Getting file from {}...", server.get_name());
        let file_source = osf!(match &server.address {
          ServerAddress::Ssh { ssh_address } => format!("{ssh_address}:"),
          ServerAddress::Localhost => "".to_string(),
        }) + server.server_directory_path.join(&file);
        ShellCmd::new("scp")
          .arg(&file_source)
          .arg(&download_directory)
          .run(&logger)
          .map_err(|e| format!("download failure: {e}"))?;

        //open file in editor
        let mut editor_command_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, download_directory.join(file_name))
          })
          .collect::<Vec<_>>();

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

        //upload file again
        ShellCmd::new("scp")
          .arg(download_directory.join(file_name))
          .arg(&file_source)
          .run(&logger)
          .map_err(|e| format!("failed to re-upload file: {e}"))?;
      }

      log!(logger, "Done!");
    }
  }

  Ok(())
}

fn check_local_file_exists<P>(path: P) -> Result<(), String>
where
  P: AsRef<Path>,
{
  let path = path.as_ref();
  if !path.is_file() {
    return Err(format!(
      "{} does not point to a file",
      path.to_string_lossy()
    ));
  }

  Ok(())
}

fn check_file_exists_on_server<P, S, D>(
  path: P,
  ssh_address: S,
  server_directory: D,
) -> Result<(), String>
where
  P: AsRef<Path>,
  S: AsRef<str>,
  D: AsRef<Path>,
{
  let full_path = server_directory.as_ref().join(path);
  match &ShellCmd::new("ssh")
    .arg(ssh_address.as_ref())
    .arg(osf!("test -f ") + &full_path)
    .collect_output()
  {
    Ok(_) => Ok(()), //file exists on file server
    Err(SpecificExecutionError {
      error: ExecutionError::BadExitStatus(_), //test failed
      ..
    }) => Err(format!(
      "File '{}' doesn't exist on file-server",
      full_path.to_string_lossy()
    )),
    Err(e) => Err(format!(
      "Failed to check whether file exists on file-server: {e}"
    )),
  }
}

fn get_home_directory() -> Result<PathBuf, String> {
  homedir::my_home()
    .map_err(|e| format!("Failed to determine home directory: {e}"))
    .and_then(|home_dir| home_dir.ok_or("Failed to find home directory".to_string()))
}

fn require_non_empty_servers<T>(servers: &[T]) -> Result<(), String> {
  require_non_empty(servers, "servers for this operation")
}

fn require_non_empty<T>(slice: &[T], slice_name: &str) -> Result<(), String> {
  if slice.is_empty() {
    Err(format!(
      "You did not provide any {slice_name}. Please see --help"
    ))?
  }
  Ok(())
}

fn start_ssh_agent(logger: &Logger) -> Result<(), String> {
  //start the ssh agent
  let agent_output = ShellCmd::new("ssh-agent")
    .arg("-s")
    .collect_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")
    .run(logger)
    .map_err(|e| format!("failed to add ssh-key: {e}"))?;
  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, RelativeLocalPathAnker::Home)
        .map_err(|e| format!("Invalid server entry '{server_entry}': {e}"))
    })
    .collect()
}

#[cfg(test)]
mod test {
  use crate::parse_server_configuration;
  use crate::server::{Server, ServerAddress};
  use std::fs;
  use std::path::PathBuf;

  #[test]
  fn test_parse_server_configuration() {
    let servers =
      parse_server_configuration("foo:bar,.:fizz/buzz").expect("valid server configuration");
    assert_eq!(
      vec![
        Server {
          address: ServerAddress::Ssh {
            ssh_address: "foo".to_string()
          },
          server_directory_path: PathBuf::from("bar"),
        },
        Server {
          address: ServerAddress::Localhost,
          server_directory_path: PathBuf::from("fizz/buzz"),
        }
      ],
      servers
    );
  }

  /// When we join an absolute path to a relative path, it becomes a relative path
  #[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);
  }

  /// When renaming a file in a folder, the folder is relevant in the new name
  #[test]
  fn rename_experiment() {
    fs::rename("test-ressources/files/test", "test-ressources/files/test1")
      .expect("failed to rename test file");
    fs::rename("test-ressources/files/test1", "test-ressources/files/test")
      .expect("failed to rename test1 file back to test");
  }

  #[test]
  fn mkdir_experiment() {
    fs::create_dir_all("./test-ressources/files/../python")
      .expect("failed to create directory with relative path");
  }
}