Moving to async

Cargo.lock

@@ -1043,7 +1043,7 @@ checksum = "436b050e76ed2903236f032a59761c1eb99e1b0aead2c257922771dab1fc8c78"
 
 [[package]]
 name = "rexec"
-version = "1.0.2"
+version = "1.0.3"
 dependencies = [
  "brace-expand",
  "clap 4.3.4",
@@ -1056,6 +1056,7 @@ dependencies = [
  "log",
  "massh",
  "regex",
+ "tokio",
  "whoami",
 ]
 

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rexec"
-version = "1.0.2"
+version = "1.0.3"
 edition = "2021"
 description = "Parallel SSH executor"
 repository = "https://github.com/house-of-vanity/rexec"
@@ -12,6 +12,7 @@ authors = ["AB <gh@hexor.ru>"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+tokio = { version = "1", features = ["full"] }
 dns-lookup = "2.0.2"
 log = "0.4.0"
 env_logger = "0.10.0"

src/main.rs

@@ -17,7 +17,6 @@ use log::{error, info};
 use massh::{MasshClient, MasshConfig, MasshHostConfig, SshAuth};
 use regex::Regex;
 
-
 // Define args
 #[derive(Parser, Debug)]
 #[command(author = "AB ab@hexor.ru", version, about = "Parallel SSH executor in Rust", long_about = None)]
@@ -28,7 +27,11 @@ struct Args {
     #[arg(short, long, help = "Use known_hosts to build servers list")]
     known_hosts: bool,
 
-    #[arg(short, long, help = "Expression to build server list. List and range expansion available. Example: 'web-[1:12]-io-{prod,dev}'")]
+    #[arg(
+        short,
+        long,
+        help = "Expression to build server list. List and range expansion available. Example: 'web-[1:12]-io-{prod,dev}'"
+    )]
     expression: String,
 
     #[arg(short, long, help = "Command to execute on servers")]
@@ -74,48 +77,41 @@ fn read_known_hosts() -> Vec<Host> {
     result
 }
 
-fn expand_string(string: String) -> Vec<Host> {
-    let mut result: Vec<String> = Vec::new();
-    let mut _result: Vec<String> = Vec::new();
+fn expand_range(start: i32, end: i32) -> Vec<String> {
+    (start..=end).map(|i| i.to_string()).collect()
+}
+
+fn expand_list(list: &str) -> Vec<String> {
+    list.split(',').map(|s| s.to_string()).collect()
+}
+
+fn expand_string(s: &str) -> Vec<Host> {
     let mut hosts: Vec<Host> = Vec::new();
+    let mut result = vec![s.to_string()];
 
-    if let Some(open_bracket_index) = string.find('[') {
-        if let Some(close_bracket_index) = string.find(']') {
-            let prefix = &string[..open_bracket_index];
-            let range = &string[open_bracket_index + 1..close_bracket_index];
-            let postfix = &string[close_bracket_index + 1..];
-
-            let parts: Vec<&str> = range.split(':').collect();
-
-            if parts.len() == 2 {
-                if let Ok(start) = parts[0].parse::<u32>() {
-                    if let Ok(end) = parts[1].parse::<u32>() {
-                        for num in start..=end {
-                            _result.push(format!("{}{}{}", prefix, num, postfix));
+    while let Some(r) = result.iter().find(|s| s.contains('[')) {
+        let r = r.clone();
+        let start = r.find('[').unwrap();
+        let end = r.find(']').unwrap();
+        let colon = r.find(':').unwrap();
+        let low = r[start+1..colon].parse::<i32>().unwrap();
+        let high = r[colon+1..end].parse::<i32>().unwrap();
+        result.retain(|s| s != &r);
+        for val in expand_range(low, high) {
+            let new_str = format!("{}{}{}", &r[..start], val, &r[end+1..]);
+            result.push(new_str);
         }
     }
-                }
-            }
-        }
-    } else {
-        _result.push(String::from(string));
-    }
 
-    for string in _result {
-        if let Some(open_brace_index) = string.find('{') {
-            if let Some(close_brace_index) = string.find('}') {
-                let prefix = &string[..open_brace_index];
-                let list = &string[open_brace_index + 1..close_brace_index];
-                let postfix = &string[close_brace_index + 1..];
-
-                let items: Vec<&str> = list.split(',').collect();
-
-                for item in items {
-                    result.push(format!("{}{}{}", prefix, item, postfix));
-                }
-            }
-        } else {
-            result.push(String::from(string));
+    while let Some(r) = result.iter().find(|s| s.contains('{')) {
+        let r = r.clone();
+        let start = r.find('{').unwrap();
+        let end = r.find('}').unwrap();
+        let list = &r[start+1..end];
+        result.retain(|s| s != &r);
+        for val in expand_list(list) {
+            let new_str = format!("{}{}{}", &r[..start], val, &r[end+1..]);
+            result.push(new_str);
         }
     }
 
@@ -128,7 +124,8 @@ fn expand_string(string: String) -> Vec<Host> {
     hosts
 }
 
-fn main() {
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
     env_logger::Builder::from_env(Env::default().default_filter_or("info"))
         .format_timestamp(None)
         .format_target(false)
@@ -153,7 +150,7 @@ fn main() {
             .collect()
     } else {
         info!("Using string expansion to build server list.");
-        expand_string(args.expression)
+        expand_string(&args.expression)
     };
 
     // Dedup hosts from known_hosts file
@@ -245,4 +242,5 @@ fn main() {
     } else {
         warn!("Stopped");
     }
+    Ok(())
 }