multi-ssh/src/main.rs

865 lines
30 KiB
Rust

mod action;
mod environment;
mod file;
#[cfg(test)]
mod integration_test;
mod logger;
mod os_str_extension;
mod os_string_builder;
mod server;
mod shell_interface;
use crate::action::{Action, FileAction, ServerActions};
use crate::environment::{Environment, Prod};
use crate::file::{FileMatcher, FileNameInfo};
use crate::logger::{LogLevel, Logger};
use crate::os_str_extension::OsStrExtension;
use crate::os_string_builder::ReplaceWithOsStr;
use crate::server::{RelativeLocalPathAnker, ServerAddress};
use crate::shell_interface::{ScpParam, ServerCommand, ShellCommand};
use clap::{Parser, Subcommand, ValueEnum};
use lazy_regex::{lazy_regex, Lazy, Regex};
use server::{Server, ServerReference};
use std::cell::LazyCell;
use std::ffi::OsString;
use std::hash::Hash;
use std::io::Write;
use std::iter::once;
use std::path::{Path, PathBuf};
use std::{env, fs, io};
const SERVERS_ENV_VAR: &str = "MSSH_SERVERS";
const EDITOR_ENV_VAR: &str = "MSSH_EDITOR";
const FILE_PLACEHOLDER: &str = "<file>";
/// 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)]
pub 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 = 0..)]
servers: Vec<String>,
/// 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<String>,
/// 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 = 'd', 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>,
/// Only upload files which are not present yet on the target server
#[arg(short, long, default_value = "false")]
pure: bool,
/// The prefixes of the names of files not to upload
#[arg(short, long, num_args = 0..)]
exclude: Vec<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,
}
#[derive(Debug, Default)]
pub struct Application<E> {
pub environment: E,
}
impl<E> Application<E>
where
E: Environment,
{
pub fn run(&mut self) -> Result<(), String> {
let args = Args::try_parse_from(self.environment.args_os()).map_err(|e| e.to_string())?;
self.run_with_args(args)
}
pub fn run_with_args(&mut self, args: Args) -> Result<(), String> {
macro_rules! env {
() => {
&mut self.environment
};
}
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(|| self.parse_server_configuration_from_env());
let servers = args
.servers
.into_iter()
.map(|ref_str| {
ServerReference::from_str(&ref_str, || self.get_home_directory())
.map_err(|e| format!("Invalid server reference '{ref_str}': {e}"))
})
.collect::<Result<Vec<_>, _>>()?
.into_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,
pure,
exclude,
} => {
Self::require_non_empty_servers(&servers)?;
Self::require_non_empty(&files, "files to upload")?;
//resolve file server
let file_server = match file_server {
Some(ref_str) => {
let server_reference =
ServerReference::from_str(&ref_str, || self.get_home_directory())
.map_err(|e| format!("Invalid file-server reference '{ref_str}': {e}"))?;
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,
};
self.start_ssh_agent(&logger)?;
//make sure files exist
match &file_server {
Some(file_server) => match &file_server.address {
ServerAddress::Ssh { ssh_address } => {
//canonicalize remote files -> also makes sure they exist
files = files
.iter()
.map(|file| {
let output = ShellCommand::Ssh {
address: ssh_address.to_string(),
server_command: ServerCommand::Realpath {
path: file_server.server_directory_path.join(file),
},
}
.in_env(env!())
.output()
.into_result()
.map_err(|e| format!("Failed to canonicalize files: {e}"))?;
if !output.status.success {
Err(format!(
"Path doesn't match any files on file-server: {}",
file.to_string_lossy()
))?;
}
let denoted_files = output
.stdout
.split(b'\n') //split at line breaks
.into_iter()
.filter(|file_name| !file_name.is_empty()) //needed since realpath sometimes gives us empty lines
.map(PathBuf::from)
.collect::<Vec<_>>();
Ok(denoted_files)
})
.collect::<Result<Vec<_>, String>>()?
.into_iter()
.flatten()
.collect();
}
ServerAddress::Localhost => files
.iter()
.map(|file| file_server.server_directory_path.join(file))
.try_for_each(Self::check_local_file_exists)?,
},
None => files.iter().try_for_each(Self::check_local_file_exists)?,
}
let file_details = files
.into_iter()
.map(|file| {
FileNameInfo::try_from(file.clone())
.map(|info| (file.clone(), info))
.map_err(|e| format!("Bad file '{}': {e}", file.to_string_lossy()))
})
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.filter(|(_, info)| {
!exclude
.iter()
.any(|exclude| info.to_full_file_name().starts_with(exclude))
})
.collect::<Vec<_>>();
log!(logger, debug, "Files to upload: ");
for (file, _) in &file_details {
log!(logger, debug, "- {}", file.to_string_lossy());
}
//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 present_file_names: Vec<OsString> = match &server.address {
ServerAddress::Ssh { ssh_address } => ShellCommand::Ssh {
address: ssh_address.to_string(),
server_command: ServerCommand::Ls {
dir: working_directory.clone(),
},
}
.in_env(env!())
.output()
.and_expect_success()
.into_result_with_error_logging(&logger)
.map_err(|e| {
format!(
"Failed to query present files on server {}: {e}",
server.get_name()
)
})?
.stdout
.split(b'\n')
.into_iter()
.map(OsString::from)
.collect(),
ServerAddress::Localhost => fs::read_dir(&working_directory)
.map_err(|e| format!("Failed to get files in working directory: {e}"))?
.map(|entry| {
entry.map_err(|e| format!("Failed to access directory entry: {e}"))
})
.collect::<Result<Vec<_>, _>>()?
.into_iter()
.filter_map(|entry| {
if entry.path().is_file() {
Some(entry.file_name())
} else {
None
}
})
.collect(),
};
file_details
.iter()
.flat_map(|(file, file_name_info)| {
let mut file_matcher = FileMatcher::from(
file_name
.as_ref()
.map(OsString::from)
.unwrap_or(file_name_info.name.to_os_string()),
);
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");
if pure && present_file_names.iter().any(|file| *file == file_name) {
log!(
logger,
debug,
"file is already present on {}: {}",
server.get_name(),
file_name.to_string_lossy()
);
return vec![]; //ignore that file, since it is already present
}
match old_version_policy {
OldVersionPolicy::Ignore => {
if !present_file_names.iter().any(|file| *file == file_name) {
vec![add_action] //file doesn't exist yet
} else {
vec![FileAction::new(file, Action::Replace).expect("path points to file")]
}
}
OldVersionPolicy::Archive => present_file_names
.iter()
.filter(|file| file_matcher.matches(file))
.map(|file| {
FileAction::new(
file,
Action::rename(
osf!(file)
+ file
.to_string_lossy()
.chars()
.last()
.unwrap_or('1')
.to_string(),
),
)
.expect("path points to file")
})
.chain(once(add_action))
.collect(),
OldVersionPolicy::Delete => {
let mut actions = present_file_names
.iter()
.filter(|file| file_matcher.matches(file))
.map(|present_file| {
//special case -> file has the same name as current file, then we just need to replace it
if *present_file == file_name {
FileAction::new(file, Action::Replace).expect("path points to file")
} else {
FileAction::new(present_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>>()?
.into_iter()
.filter(|server_actions| !server_actions.actions.is_empty())
.collect::<Vec<_>>();
if actions.is_empty() {
log!(logger, "Nothing to be done, everything is up to date");
return Ok(());
}
log!(logger, "The following actions will be performed: ");
for server_actions in &actions {
log!(logger, "{server_actions}");
log!(logger, debug, "Detailed file actions: ");
for file_action in &server_actions.actions {
log!(logger, debug, "{file_action:?}");
}
}
if !no_confirm && !self.confirm("Continue?", true) {
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 source = match &file_server {
Some(file_server) => ScpParam::from((
file_server,
file_server.server_directory_path.join(&file_action.file),
)),
None => ScpParam::from(file_action.file.as_path()),
};
let destination = ScpParam::from((server, &server_actions.working_directory));
ShellCommand::Scp {
source,
destination,
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&logger)
.map_err(|e| format!("upload failure: {e}"))?;
}
Action::Delete => match &server.address {
ServerAddress::Ssh { ssh_address } => {
ShellCommand::Ssh {
address: ssh_address.to_string(),
server_command: ServerCommand::Rm {
file: server_actions.working_directory.join(&file_action.file),
},
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&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 } => {
ShellCommand::Ssh {
address: ssh_address.to_string(),
server_command: ServerCommand::Mv {
source: server_actions.working_directory.join(&file_action.file),
destination: server_actions.working_directory.join(&new_name),
},
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&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 } => {
self.start_ssh_agent(&logger)?;
Self::require_non_empty_servers(&servers)?;
for server in servers {
log!(logger, "Running command on '{}'...", server.get_name());
match &server.address {
ServerAddress::Ssh { ssh_address } => {
ShellCommand::Ssh {
address: ssh_address.to_string(),
server_command: ServerCommand::Execute {
working_directory: server.server_directory_path.clone(),
command: OsString::from(&command),
},
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&logger)
.map_err(|e| format!("{e}"))?;
}
ServerAddress::Localhost => {
let command = shell_words::split(&command)
.map_err(|e| format!("failed to parse command: {e}"))?
.into_iter()
.map(OsString::from)
.collect();
ShellCommand::Execute {
working_directory: server.server_directory_path.clone(),
command,
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&logger)?;
}
}
}
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 = self
.get_home_directory()
.map_err(|e| format!("Missing 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
&& self.confirm(
format!("{duplication_notification}. Do you want to replace it?"),
false,
)
{
break 'duplicate_check;
}
return Err(format!(
"{duplication_notification}. You can override it with --override or -f"
));
}
}
Self::require_non_empty_servers(&servers)?;
self.start_ssh_agent(&logger)?;
for server in servers {
log!(logger, "Getting file from {}...", server.get_name());
let source = ScpParam::from((&server, server.server_directory_path.join(&file)));
ShellCommand::Scp {
source: source.clone(),
destination: ScpParam::from(download_directory.as_path()),
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&logger)
.map_err(|e| format!("download failure: {e}"))?;
//open file in editor
let editor_command = 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<_>>();
ShellCommand::Editor(editor_command)
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&logger)
.map_err(|e| format!("failed to open file in editor: {e}"))?;
//upload file again
ShellCommand::Scp {
source: ScpParam::from(download_directory.join(file_name).as_path()),
destination: source,
}
.in_env(env!())
.run_logged(&logger)
.and_expect_success()
.into_result_with_error_logging(&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 require_non_empty_servers<T>(servers: &[T]) -> Result<(), String> {
Self::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(&mut self, logger: &Logger) -> Result<(), String> {
let env = &mut self.environment;
//start the ssh agent
let agent_output = ShellCommand::SshAgent
.in_env(env)
.output()
.and_expect_success()
.into_result_with_error_logging(logger)
.map_err(|e| format!("Failed to start ssh agent: {e}"))?;
let agent_stdout = &agent_output
.stdout
.into_string()
.map_err(|_| "ssh-agent returned invalid utf-8 - how did this even happen?")?;
//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
ShellCommand::ShhAdd
.in_env(env)
.run_logged(logger)
.and_expect_success()
.into_result_with_error_logging(logger)
.map_err(|e| format!("failed to add ssh-key: {e}"))?;
Ok(())
}
fn parse_server_configuration_from_env(&self) -> Result<Vec<Server>, String> {
self
.environment
.var(SERVERS_ENV_VAR)
.map_err(|_| format!("Missing environment variable {}", SERVERS_ENV_VAR))
.and_then(|value| parse_server_configuration(&value, || self.get_home_directory()))
}
fn get_home_directory(&self) -> Result<PathBuf, String> {
self
.environment
.get_home_directory()
.ok_or("Failed to find your home directory".to_string())
}
fn confirm<S>(&mut self, prompt: S, default_value: bool) -> bool
where
S: ToString,
{
loop {
print!(
"{}[{}]",
prompt.to_string(),
if default_value { "Y|n" } else { "y|N" }
);
io::stdout().flush().expect("failed to flush stdout");
let line = self
.environment
.read_line()
.expect("Failed to read console input");
match line.to_lowercase().as_str() {
"" => return default_value,
"y" | "yes" => return true,
"n" | "no" => return false,
_ => println!("Invalid input, please choose one of the provided options"),
}
}
}
}
fn main() -> Result<(), String> {
Application::<Prod>::default().run()
}
fn parse_server_configuration<F>(
config_str: &str,
get_home_directory: F,
) -> Result<Vec<Server>, String>
where
F: Fn() -> Result<PathBuf, String>,
{
config_str
.split(',')
.map(|server_entry| {
Server::from_str(
server_entry,
RelativeLocalPathAnker::Home,
&get_home_directory,
)
.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() {
//setup directory structure for test
let home_dir = PathBuf::from("target/test");
const LOCAL_SERVER_DIR: &str = "fizz/buzz";
fs::create_dir_all(home_dir.join(LOCAL_SERVER_DIR)).expect("failed to create server directory");
let servers = parse_server_configuration(&format!("foo:bar,.:{LOCAL_SERVER_DIR}"), || {
Ok(home_dir.clone())
})
.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: home_dir
.join(LOCAL_SERVER_DIR)
.canonicalize()
.expect("home dir exists"),
}
],
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");
}
}