diff --git a/Cargo.lock b/Cargo.lock index 76c2909..9a9ab9c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "aho-corasick" @@ -244,6 +244,31 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "crypto-common" version = "0.1.6" @@ -1029,6 +1054,26 @@ dependencies = [ "getrandom", ] +[[package]] +name = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.2.16" @@ -1079,6 +1124,7 @@ dependencies = [ "massh", "openssl", "question", + "rayon", "regex", "whoami", ] diff --git a/Cargo.toml b/Cargo.toml index a024855..f209f40 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rexec" -version = "1.1.0" +version = "1.2.0" readme = "https://github.com/house-of-vanity/rexec#readme" edition = "2021" description = "Parallel SSH executor" @@ -14,6 +14,7 @@ authors = ["AB "] [dependencies] dns-lookup = "2.0.2" +rayon = "1.10" log = "0.4.0" env_logger = "0.10.0" massh = "0.6.3" diff --git a/src/main.rs b/src/main.rs index b189dcd..d5bbb7d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,6 +6,7 @@ use std::fs::read_to_string; use std::hash::Hash; use std::net::IpAddr; use std::process; +use std::sync::{Arc, Mutex}; use clap::Parser; use colored::*; @@ -15,6 +16,7 @@ use itertools::Itertools; use log::{error, info}; use massh::{MasshClient, MasshConfig, MasshHostConfig, SshAuth}; use question::{Answer, Question}; +use rayon::prelude::*; use regex::Regex; // Define args @@ -187,8 +189,6 @@ fn main() { let matched_hosts: Vec<_> = hosts.into_iter().unique().collect(); // Build MasshHostConfig hostnames list - let mut massh_hosts: Vec = vec![]; - let mut hosts_and_ips: HashMap = HashMap::new(); if args.parallel != 100 { warn!("Parallelism: {} thread{}", &args.parallel, { if args.parallel != 1 { @@ -200,23 +200,42 @@ fn main() { } info!("Matched hosts:"); - for host in matched_hosts.iter() { - let ip = match lookup_host(&host.name) { - Ok(ip) => ip[0], - Err(_) => { - error!("{} couldn't be resolved.", &host.name.red()); - continue; + let resolved_ips = Arc::new(Mutex::new(Vec::<(String, IpAddr)>::new())); + + matched_hosts.par_iter().for_each(|host| { + match lookup_host(&host.name) { + Ok(ips) if !ips.is_empty() => { + let ip = ips[0]; + + info!("{} [{}]", &host.name, ip); + + let mut results = resolved_ips.lock().unwrap(); + results.push((host.name.clone(), ip)); } - }; - info!("{} [{}]", &host.name, ip); - hosts_and_ips.insert(ip, host.name.clone()); - massh_hosts.push(MasshHostConfig { - addr: ip, - auth: None, - port: None, - user: None, - }) + Ok(_) => { + error!("DNS resolved, but IP not found: {}", &host.name.red()); + } + Err(_) => { + error!("DNS resolve failed: {}", &host.name.red()); + } + } + }); + + let mut hosts_and_ips: HashMap = HashMap::new(); + let mut massh_hosts: Vec = Vec::new(); + + if let Ok(results) = resolved_ips.lock() { + for (hostname, ip) in results.iter() { + hosts_and_ips.insert(*ip, hostname.clone()); + massh_hosts.push(MasshHostConfig { + addr: *ip, + auth: None, + port: None, + user: None, + }); + } } + // Build MasshConfig using massh_hosts vector let config = MasshConfig { default_auth: SshAuth::Agent,