From bd4c466111f1f3322edbab75668ea26cb608277b Mon Sep 17 00:00:00 2001
From: Alexandr Bogomyakov <abogomyakov@iponweb.net>
Date: Fri, 16 Jun 2023 17:38:56 +0300
Subject: [PATCH] Added dedup of known_hosts

---
 Cargo.lock  | 16 ++++++++++++++
 Cargo.toml  |  1 +
 src/main.rs | 61 ++++++++++++++++++++++++++++++++---------------------
 3 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7f9a94e..0c3c290 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -289,6 +289,12 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "either"
+version = "1.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91"
+
 [[package]]
 name = "encode_unicode"
 version = "0.3.6"
@@ -630,6 +636,15 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "itertools"
+version = "0.10.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.6"
@@ -1029,6 +1044,7 @@ dependencies = [
  "dialoguer",
  "dns-lookup",
  "env_logger",
+ "itertools",
  "lazy-st",
  "log",
  "massh",
diff --git a/Cargo.toml b/Cargo.toml
index f4cfa16..39f5a95 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,3 +16,4 @@ regex = "1"
 clap = { version = "4.3.4", features = ["derive"] }
 colored = "2.0.0"
 dialoguer = "0.10.4"
+itertools = "0.10.5"
diff --git a/src/main.rs b/src/main.rs
index b0ceb03..00cccb0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,6 +2,7 @@ use colored::*;
 use dialoguer::Confirm;
 use dns_lookup::lookup_host;
 use env_logger::Env;
+use itertools::Itertools;
 use log::{error, info};
 use massh::{MasshClient, MasshConfig, MasshHostConfig, SshAuth};
 use regex::Regex;
@@ -13,33 +14,40 @@ extern crate log;
 
 use clap::Parser;
 
-/// Simple program to greet a person
+// parse args
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
-    /// Name of the person to greet
     #[arg(short, long, default_value_t = whoami::username())]
     username: String,
 
-    #[arg(short, long)]
-    pattern: String,
+    #[arg(short, long, help = "Use known_hosts to build servers list.")]
+    expression: String,
 
-    #[arg(short, long)]
+    #[arg(short, long, help = "Command to execute on servers.")]
     command: String,
+
+    #[arg(short, long, default_value_t = false, help = "Show exit code ONLY")]
+    code: bool,
+
+    #[arg(short, long, default_value_t = 100)]
+    parallel: i32,
 }
 
-#[derive(Debug, Default)]
+// Represent line from known_hosts file
+#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
 struct KnownHost {
     name: String,
     ip: Option<IpAddr>,
 }
 
-impl KnownHost {
-    fn new(name: String) -> KnownHost {
-        KnownHost { name, ip: None }
-    }
-}
+// impl KnownHost {
+//     fn new(name: String) -> KnownHost {
+//         KnownHost { name, ip: None }
+//     }
+// }
 
+// Read known_hosts file
 fn read_known_hosts() -> Vec<KnownHost> {
     let mut result: Vec<KnownHost> = Vec::new();
 
@@ -49,10 +57,6 @@ fn read_known_hosts() -> Vec<KnownHost> {
     {
         let line = line.split(" ").collect::<Vec<&str>>();
         let hostname = line[0];
-        // let ip = match lookup_host(hostname) {
-        //     Ok(ip) => ip[0],
-        //     Err(_) => continue
-        // };
         result.push(KnownHost {
             name: hostname.to_string(),
             ip: None,
@@ -69,13 +73,18 @@ fn main() {
     let args = Args::parse();
 
     let known_hosts = read_known_hosts();
-    let re = Regex::new(&args.pattern).unwrap();
+    // Build regex
+    let re = Regex::new(&args.expression).unwrap();
+    // match hostnames from known_hosts to regex
     let mut matched_hosts: Vec<KnownHost> = known_hosts
         .into_iter()
         .filter(|r| re.is_match(&r.name.clone()))
-        //.filter(|a| a.access.iter().any(|b| b.users.contains(&user)))
         .collect();
 
+    // Dedup hosts from known_hosts file
+    let mut matched_hosts: Vec<_> = matched_hosts.into_iter().unique().collect();
+
+    // Build MasshHostConfig hostnames list
     let mut massh_hosts: Vec<MasshHostConfig> = vec![];
     info!("Matched hosts:");
     for host in matched_hosts.iter() {
@@ -95,19 +104,19 @@ fn main() {
         })
     }
 
+    // Build MasshConfig using massh_hosts vector
     let config = MasshConfig {
         default_auth: SshAuth::Agent,
         default_port: 22,
         //default_user: whoami::username(),
         default_user: "abogomyakov".to_string(),
-        threads: 10,
+        threads: args.parallel as u64,
         timeout: 0,
         hosts: massh_hosts,
     };
     let massh = MasshClient::from(&config);
 
-    // Receive the result of the command for each host and print its output.
-
+    // Ask for confirmation
     if Confirm::new()
         .with_prompt(format!(
             "Continue on following {} servers?",
@@ -118,7 +127,10 @@ fn main() {
     {
         info!("\n");
         info!("Run command on {} servers.", &config.hosts.len());
-        info!("\n"); // Run a command on all the configured hosts.
+        info!("\n");
+
+        // Run a command on all the configured hosts.
+        // Receive the result of the command for each host and print its output.
         let rx = massh.execute(args.command);
 
         while let Ok((host, result)) = rx.recv() {
@@ -129,9 +141,10 @@ fn main() {
             } else {
                 info!("Code {}", output.exit_status.to_string().red());
             };
-
-            info!("STDOUT:\n{}", String::from_utf8(output.stdout).unwrap());
-            info!("STDERR:\n{}", String::from_utf8(output.stderr).unwrap());
+            if !args.code {
+                info!("STDOUT:\n{}", String::from_utf8(output.stdout).unwrap());
+                info!("STDERR:\n{}", String::from_utf8(output.stderr).unwrap());
+            }
         }
     } else {
         warn!("Stopped");