diff --git a/Cargo.lock b/Cargo.lock index 0c3c290..7c53800 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -119,6 +119,12 @@ dependencies = [ "generic-array", ] +[[package]] +name = "brace-expand" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3adb80ee272c844254166ea32c8ae11c211b3639a293fdde41b1645b6be2c62" + [[package]] name = "bumpalo" version = "3.13.0" @@ -1039,6 +1045,7 @@ checksum = "436b050e76ed2903236f032a59761c1eb99e1b0aead2c257922771dab1fc8c78" name = "rexec" version = "0.1.0" dependencies = [ + "brace-expand", "clap 4.3.4", "colored", "dialoguer", diff --git a/Cargo.toml b/Cargo.toml index 39f5a95..d8efcaa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,3 +17,4 @@ clap = { version = "4.3.4", features = ["derive"] } colored = "2.0.0" dialoguer = "0.10.4" itertools = "0.10.5" +brace-expand = "0.1.0" diff --git a/src/main.rs b/src/main.rs index d2309da..e94a06c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,4 @@ +use brace_expand::brace_expand; use colored::*; use dialoguer::Confirm; use dns_lookup::lookup_host; @@ -5,27 +6,31 @@ use env_logger::Env; use itertools::Itertools; use log::{error, info}; use massh::{MasshClient, MasshConfig, MasshHostConfig, SshAuth}; -use regex::Regex; +use regex::{Error, Regex}; use std::collections::{HashMap, HashSet}; use std::fs::read_to_string; use std::hash::Hash; use std::net::IpAddr; use std::str::FromStr; +use std::process; #[macro_use] extern crate log; use clap::Parser; -// parse args +// Define args #[derive(Parser, Debug)] -#[command(author, version, about, long_about = None)] +#[command(author = "AB ab@hexor.ru", version, about = "Parallel SSH executor in Rust", long_about = None)] struct Args { #[arg(short, long, default_value_t = whoami::username())] username: String, - #[arg(short, long = "kh", help = "Use known_hosts to build servers list")] - known_hosts: String, + #[arg(short, long, help = "Use known_hosts to build servers list")] + known_hosts: bool, + + #[arg(short, long, help = "Expression to build server list")] + expression: String, #[arg(short, long, help = "Command to execute on servers")] command: String, @@ -33,7 +38,12 @@ struct Args { #[arg(long, default_value_t = false, help = "Show exit code ONLY")] code: bool, - #[arg(long, default_value_t = false, help = "Don't ask for confirmation")] + #[arg( + short = 'f', + long, + default_value_t = false, + help = "Don't ask for confirmation" + )] noconfirm: bool, #[arg(short, long, default_value_t = 100)] @@ -42,20 +52,14 @@ struct Args { // Represent line from known_hosts file #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] -struct KnownHost { +struct Host { name: String, ip: Option, } -// impl KnownHost { -// fn new(name: String) -> KnownHost { -// KnownHost { name, ip: None } -// } -// } - // Read known_hosts file -fn read_known_hosts() -> Vec { - let mut result: Vec = Vec::new(); +fn read_known_hosts() -> Vec { + let mut result: Vec = Vec::new(); for line in read_to_string(format!("/home/{}/.ssh/known_hosts", whoami::username())) .unwrap() @@ -63,7 +67,19 @@ fn read_known_hosts() -> Vec { { let line = line.split(" ").collect::>(); let hostname = line[0]; - result.push(KnownHost { + result.push(Host { + name: hostname.to_string(), + ip: None, + }) + } + result +} + +fn expand_expression(expr: &str) -> Vec { + let mut result: Vec = Vec::new(); + + for hostname in brace_expand(expr) { + result.push(Host { name: hostname.to_string(), ip: None, }) @@ -78,17 +94,29 @@ fn main() { .init(); let args = Args::parse(); - let known_hosts = read_known_hosts(); - // Build regex - let re = Regex::new(&args.known_hosts).unwrap(); - // match hostnames from known_hosts to regex - let mut matched_hosts: Vec = known_hosts - .into_iter() - .filter(|r| re.is_match(&r.name.clone())) - .collect(); + let hosts = if args.known_hosts { + info!("Using ~/.ssh/known_hosts to build server list."); + let known_hosts = read_known_hosts(); + // Build regex + let re = match Regex::new(&args.expression) { + Ok(result) => result, + Err(e) => { + error!("Error parsing regex. {}", e); + process::exit(1); + } + }; + // match hostnames from known_hosts to regex + known_hosts + .into_iter() + .filter(|r| re.is_match(&r.name.clone())) + .collect() + } else { + info!("Using string expansion to build server list."); + expand_expression(&args.expression) + }; // Dedup hosts from known_hosts file - let mut matched_hosts: Vec<_> = matched_hosts.into_iter().unique().collect(); + let mut matched_hosts: Vec<_> = hosts.into_iter().unique().collect(); // Build MasshHostConfig hostnames list let mut massh_hosts: Vec = vec![]; @@ -98,7 +126,7 @@ fn main() { let ip = match lookup_host(&host.name) { Ok(ip) => ip[0], Err(_) => { - error!("{} couldn't ve resolved.", &host.name.red()); + error!("{} couldn't be resolved.", &host.name.red()); continue; } }; @@ -143,11 +171,23 @@ fn main() { while let Ok((host, result)) = rx.recv() { let ip: String = host.split('@').collect::>()[1] .split(':') - .collect::>()[0].to_string(); + .collect::>()[0] + .to_string(); let ip = ip.parse::().unwrap(); //let ip = hosts_and_ips.get(&ip).unwrap_or(&"Couldn't parse IP".yellow().bold().to_string()); - info!("{}", hosts_and_ips.get(&ip).unwrap_or(&"Couldn't parse IP".yellow().bold().to_string())); - let output = result.unwrap(); + info!( + "{}", + hosts_and_ips + .get(&ip) + .unwrap_or(&"Couldn't parse IP".yellow().bold().to_string()) + ); + let output = match result { + Ok(output) => output, + Err(e) => { + error!("Can't access server: {}", e); + continue + } + }; if output.exit_status == 0 { info!("Code {}", output.exit_status.to_string().green()); } else {