1 Commits

Author SHA1 Message Date
c699bf1849 Format Rust code using rustfmt 2023-07-03 12:54:56 +00:00
9 changed files with 1640 additions and 950 deletions

View File

@ -1,10 +0,0 @@
pkgbase = rexec
pkgdesc = Parallel SSH executor
pkgver = 1.0.6
pkgrel = 1
url = https://github.com/house-of-vanity/rexec
arch = x86_64
license = WTFPL
makedepends = cargo
pkgname = rexec

View File

@ -1,140 +1,80 @@
name: Rust static build and publish
on:
push:
tags:
- 'v*.*.*'
- '*'
env:
CARGO_TERM_COLOR: always
BUILD_TARGET: x86_64-unknown-linux-musl
BINARY_NAME: rexec
jobs:
build:
name: Build static binary
runs-on: ${{ matrix.os }}
strategy:
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
include:
- os: ubuntu-latest
build_target: x86_64-unknown-linux-musl
platform_name: linux-amd64
- os: windows-latest
build_target: x86_64-pc-windows-msvc
platform_name: windows-amd64
- os: macos-latest
build_target: aarch64-apple-darwin
platform_name: macos-arm64
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- uses: actions/checkout@v4
- name: Cache Cargo registry
uses: actions/cache@v4
with:
path: ~/.cargo/registry
key: ${{ runner.os }}-cargo-registry-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-registry-
- name: Cache Cargo index
uses: actions/cache@v4
with:
path: ~/.cargo/git
key: ${{ runner.os }}-cargo-index-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-index-
- name: Cache Cargo build
uses: actions/cache@v4
with:
path: target
key: ${{ runner.os }}-cargo-build-${{ hashFiles('**/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-build-
- uses: actions/checkout@v3
- uses: dtolnay/rust-toolchain@stable
- uses: mbrobbel/rustfmt-check@master
with:
token: ${{ secrets.GITHUB_TOKEN }}
- name: Install rust targets
run: rustup target add ${{ matrix.build_target }}
- name: Build Linux MUSL
if: matrix.os == 'ubuntu-latest'
- name: Build-musl
uses: gmiam/rust-musl-action@master
with:
args: cargo build --target ${{ matrix.build_target }} --release
- name: Build MacOS
if: matrix.os == 'macos-latest'
run: cargo build --target ${{ matrix.build_target }} --release
- name: Build Windows
if: matrix.os == 'windows-latest'
run: cargo build --target ${{ matrix.build_target }} --release
- name: Upload artifact
uses: actions/upload-artifact@v4
args: cargo build --target $BUILD_TARGET --release
- uses: actions/upload-artifact@v2
name: Upload artifact
with:
name: ${{ env.BINARY_NAME }}_${{ matrix.platform_name }}
path: target/${{ matrix.build_target }}/release/${{ env.BINARY_NAME }}*
name: ${{ env.BINARY_NAME }}_x86_64-linux-musl
path: target/x86_64-unknown-linux-musl/release/${{ env.BINARY_NAME }}
# build_arch:
# name: Arch Linux package
# runs-on: ubuntu-latest
# permissions:
# contents: write
# steps:
# - uses: actions/checkout@v3
# - name: Install latest nightly
# uses: actions-rs/toolchain@v1
# with:
# toolchain: stable
# override: true
# target: x86_64-unknown-linux-musl
# components: rustfmt
# - name: Install cargo-aur
# run: cargo install cargo-aur
# - name: Build PKGBUILD
# uses: actions-rs/cargo@v1
# with:
# command: aur
# - uses: actions/upload-artifact@v2
# name: Upload artifact PKGBUILD
# with:
# path: PKGBUILD
# name: PKGBUILD
release:
name: Create Release Page
needs: build
name: Release binary
needs: ["build"]
runs-on: ubuntu-latest
outputs:
upload_url: ${{ steps.create_release.outputs.upload_url }}
permissions:
contents: write
steps:
- uses: actions/checkout@v4
- name: Get version
run: echo "VERSION=$(grep -P '^version = \"\d+\.\d+\.\d+\"' Cargo.toml | awk -F '\"' '{print $2}')" >> $GITHUB_ENV
- name: Show version
run: echo ${{ env.VERSION }}
- uses: actions/download-artifact@master
id: download_binary
with:
name: ${{ env.BINARY_NAME }}_x86_64-linux-musl
path: ${{ env.BINARY_NAME }}_x86_64-linux-musl
# - uses: actions/download-artifact@master
# id: download_pkgbuild
# with:
# name: PKGBUILD
# path: PKGBUILD
- uses: ncipollo/release-action@v1
id: create_release
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
name: Update release page
with:
artifacts: "${{ steps.download_binary.outputs.download-path }}/*"
allowUpdates: true
body: "Release rexec ${{ env.VERSION }}. Static build for Linux (x86_64), Windows (x86_64) и MacOS (arm64)."
upload:
name: Upload Release Assets
needs: release
runs-on: ${{ matrix.os }}
strategy:
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
include:
- os: ubuntu-latest
platform_name: linux-amd64
- os: windows-latest
platform_name: windows-amd64
- os: macos-latest
platform_name: macos-arm64
steps:
- uses: actions/checkout@v4
- uses: actions/download-artifact@v4
name: Download ${{ matrix.platform_name }} artifact
with:
name: ${{ env.BINARY_NAME }}_${{ matrix.platform_name }}
path: ${{ env.BINARY_NAME }}_${{ matrix.platform_name }}
- name: Upload Release Asset
uses: actions/upload-release-asset@v1
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
with:
upload_url: ${{ needs.release.outputs.upload_url }}
asset_path: ${{ env.BINARY_NAME }}_${{ matrix.platform_name }}/${{ env.BINARY_NAME }}${{ matrix.platform_name == 'windows-amd64' && '.exe' || '' }}
asset_name: ${{ env.BINARY_NAME }}_${{ matrix.platform_name }}${{ matrix.platform_name == 'windows-amd64' && '.exe' || '' }}
asset_content_type: application/octet-stream
body: Static build for x86_64-linux using MUSL

6
.gitignore vendored
View File

@ -1,7 +1 @@
target
pkg
*zst
.rustup
.cargo
.viminfo
.bash_history

View File

@ -6,26 +6,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [1.1.0] - 2023-09-13
### Improved
- Logging. Fix bugs
## [1.0.7] - 2023-08-31
### Changed
- Ci
## [1.0.4] - 2023-08-14
### Changed
- Expansion host strings.
## [1.0.1] - 2023-06-20
### Changed

1607
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -1,7 +1,6 @@
[package]
name = "rexec"
version = "1.5.1"
readme = "https://github.com/house-of-vanity/rexec#readme"
version = "1.0.2"
edition = "2021"
description = "Parallel SSH executor"
repository = "https://github.com/house-of-vanity/rexec"
@ -13,16 +12,15 @@ authors = ["AB <gh@hexor.ru>"]
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies]
dns-lookup = "2"
rayon = "1.10"
log = "0.4"
env_logger = "0.11"
whoami = "1"
lazy-st = "1"
dns-lookup = "2.0.2"
log = "0.4.0"
env_logger = "0.10.0"
massh = "0.6.3"
whoami = "1.4.0"
lazy-st = "0.2.2"
regex = "1"
clap = { version = "4.5", features = ["derive"] }
colored = "3"
itertools = "0.14"
clap = { version = "4.3.4", features = ["derive"] }
colored = "2.0.0"
dialoguer = "0.10.4"
itertools = "0.10.5"
brace-expand = "0.1.0"
question = "0.2.2"
lazy_static = "1.5.0"

View File

@ -1,14 +0,0 @@
# Maintainer: AB <gh@hexor.ru>
pkgname=rexec
pkgver=some_ver
pkgrel=1
pkgdesc="Parallel SSH executor"
url="https://github.com/house-of-vanity/rexec"
license=("WTFPL")
arch=("x86_64")
package() {
install -Dm755 "$startdir/target/x86_64-unknown-linux-musl/release/rexec" "$pkgdir/usr/bin/rexec"
}

View File

@ -1,20 +1,14 @@
# rexec
**Parallel SSH executor** in Rust with static binary. It can read servers from `~/.ssh/known_hosts`
Parallel SSH executor in Rust. It can read servers from `~/.ssh/known_hosts`
or even expand servers from bash expanshion line `cassandra.{prod,dev}.example.com`
or even expand servers from bash expanshion line `cassandra-[1:5].{prod,dev}.example.com`
[![Rust-static-build](https://github.com/house-of-vanity/rexec/actions/workflows/release.yml/badge.svg)](https://github.com/house-of-vanity/rexec/actions/workflows/release.yml)
---
## Usage
```shell
_|_|_|_| _| _| _|_|_|_| _|_|_|
_| _|_| _| _| _| _| _|
_|_| _|_|_| _| _|_|_| _|
_| _| _| _| _| _|
_| _|_|_|_| _| _| _|_|_|_| _|_|_|
[ab@test]$ ./rexec -k admin-gce-sc.* --help
Parallel SSH executor in Rust
Usage: rexec [OPTIONS] --expression <EXPRESSION> --command <COMMAND>
@ -30,33 +24,39 @@ Options:
-h, --help Print help
-V, --version Print version
```
---
![image](https://github.com/house-of-vanity/rexec/assets/4666566/4c52915d-2bc1-46b9-9833-b0d7c0527f2d)
## Examples
```shell
$ rexec -f \
-e 'cassandra-gce-or-[1:2]' \
-u ab \
-c 'uname -r; date'
[INFO ] Using string expansion to build server list.
[ab@test debug]$ ./rexec -u ab -k -c uptime -e admin.* -f
[INFO ] Matched hosts:
[INFO ] cassandra-gce-or-1.prod.example.com [2.22.123.79]
[INFO ] cassandra-gce-or-2.prod.example.com [2.22.123.158]
Continue on following 2 servers? y
[INFO ] Run command on 2 servers.
[INFO ] admin-gce-sc-1.lca-prod.** [35.211.27.195]
[INFO ] admin-gce-sc-1.mmk-prod.** [35.211.79.202]
[ERROR] admin-gce-sc-1.led-prod.** couldn't be resolved.
[INFO ] admin-gce-sc-1.msq-dev.** [35.211.0.24]
[ERROR] admin-gce-sc-1.hui-dev.** couldn't be resolved.
cassandra-gce-or-1.prod.example.com
Exit code [0] / stdout 45 bytes / stderr 0 bytes
STDOUT
║ 5.15.0-1040-gcp
║ Thu Sep 7 13:44:40 UTC 2023
[INFO ] Run command on 3 servers.
cassandra-gce-or-2.prod.example.com
Exit code [0] / stdout 45 bytes / stderr 0 bytes
STDOUT
║ 5.15.0-1040-gcp
║ Thu Sep 7 13:44:40 UTC 2023
[INFO ] admin.gnb-prod.**
Code 0
STDOUT:
23:31:21 up 294 days, 14:14, 0 users, load average: 0.53, 0.64, 0.52
STDERR:
[INFO ] admin.abe-prod.**
Code 0
STDOUT:
23:31:22 up 154 days, 9:24, 0 users, load average: 0.31, 0.25, 0.18
STDERR:
[INFO ] admin-gce-be-1.toy-prod.**
Code 0
STDOUT:
23:31:22 up 98 days, 6:20, 0 users, load average: 0.88, 0.74, 0.80
STDERR:
no server running on /tmp/tmux-1019/default
open terminal failed: not a terminal
```

View File

@ -1,64 +1,45 @@
#[macro_use]
extern crate log;
use std::collections::HashMap;
use std::fs::read_to_string;
use std::hash::Hash;
use std::io::{BufRead, BufReader};
use std::net::IpAddr;
use std::process::{self, Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread;
use std::process;
use clap::Parser;
use colored::*;
use dialoguer::Confirm;
use dns_lookup::lookup_host;
use env_logger::Env;
use itertools::Itertools;
use lazy_static::lazy_static;
use log::{error, info, warn};
use question::{Answer, Question};
use rayon::prelude::*;
use log::{error, info};
use massh::{MasshClient, MasshConfig, MasshHostConfig, SshAuth};
use regex::Regex;
// Global state to track the currently open block
lazy_static! {
static ref CURRENT_BLOCK: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
}
// Define command-line arguments using the clap library
// Define args
#[derive(Parser, Debug)]
#[command(author = "AB ab@hexor.ru", version, about = "Parallel SSH executor in Rust", long_about = None)]
struct Args {
/// Username for SSH connections (defaults to current system user)
#[arg(short = 'u', short_alias = 'l', long, default_value_t = whoami::username())]
#[arg(short, long, default_value_t = whoami::username())]
username: String,
/// Flag to use known_hosts file for server discovery instead of pattern expansion
#[arg(
short,
long,
help = "Use known_hosts to build servers list instead of string expansion."
)]
#[arg(short, long, help = "Use known_hosts to build servers list")]
known_hosts: bool,
/// Server name patterns with expansion syntax
/// Examples: 'web-[1:12]-io-{prod,dev}' expands to multiple servers
#[arg(
short,
long,
num_args = 1..,
help = "Expression to build server list. List and range expansion are supported. Example: 'web-[1:12]-io-{prod,dev}'"
help = "Expression to build server list. List and range expansion available. Example: 'web-[1:12]-io-{prod,dev}'"
)]
expression: Vec<String>,
expression: String,
/// Command to execute on each server
#[arg(short, long, help = "Command to execute on servers")]
command: String,
/// Display only exit codes without command output
#[arg(long, default_value_t = false, help = "Show exit code ONLY")]
code: bool,
/// Skip confirmation prompt before executing commands
#[arg(
short = 'f',
long,
@ -67,104 +48,21 @@ struct Args {
)]
noconfirm: bool,
/// Maximum number of parallel SSH connections
#[arg(short, long, default_value_t = 100)]
parallel: i32,
}
/// Host representation for both known_hosts entries and expanded patterns
// Represent line from known_hosts file
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
struct Host {
/// Hostname or IP address as a string
name: String,
/// Resolved IP address (if available)
ip: Option<IpAddr>,
}
/// Find common domain suffix across all hostnames to simplify output display
///
/// This function analyzes all hostnames to identify a common domain suffix
/// which can be shortened during display to improve readability.
///
/// # Arguments
/// * `hostnames` - A slice of strings containing all server hostnames
///
/// # Returns
/// * `Option<String>` - The common suffix if found, or None
fn find_common_suffix(hostnames: &[String]) -> Option<String> {
if hostnames.is_empty() {
return None;
}
// Don't truncate if only one host
if hostnames.len() == 1 {
return None;
}
let first = &hostnames[0];
// Start with assumption that the entire first hostname is the common suffix
let mut common = first.clone();
// Iterate through remaining hostnames, reducing the common part
for hostname in hostnames.iter().skip(1) {
// Exit early if no common part remains
if common.is_empty() {
return None;
}
// Find common suffix with current hostname
let mut new_common = String::new();
// Search for common suffix by comparing characters from right to left
let mut common_chars = common.chars().rev();
let mut hostname_chars = hostname.chars().rev();
loop {
match (common_chars.next(), hostname_chars.next()) {
(Some(c1), Some(c2)) if c1 == c2 => new_common.insert(0, c1),
_ => break,
}
}
common = new_common;
}
// Ensure the common part is a valid domain suffix (starts with a dot)
if common.is_empty() || !common.starts_with('.') {
return None;
}
// Return the identified common suffix
Some(common)
}
/// Shorten hostname by removing the common suffix and replacing with an asterisk
///
/// # Arguments
/// * `hostname` - The original hostname
/// * `common_suffix` - Optional common suffix to remove
///
/// # Returns
/// * `String` - Shortened hostname or original if no common suffix
fn shorten_hostname(hostname: &str, common_suffix: &Option<String>) -> String {
match common_suffix {
Some(suffix) if hostname.ends_with(suffix) => {
let short_name = hostname[..hostname.len() - suffix.len()].to_string();
format!("{}{}", short_name, "*")
}
_ => hostname.to_string(),
}
}
/// Read and parse the SSH known_hosts file to extract server names
///
/// # Returns
/// * `Vec<Host>` - List of hosts found in the known_hosts file
// Read known_hosts file
fn read_known_hosts() -> Vec<Host> {
let mut result: Vec<Host> = Vec::new();
// Read known_hosts file from the user's home directory
for line in read_to_string(format!("/home/{}/.ssh/known_hosts", whoami::username()))
.unwrap()
.lines()
@ -179,91 +77,51 @@ fn read_known_hosts() -> Vec<Host> {
result
}
/// Expand a numeric range in the format [start:end] to a list of strings
///
/// # Arguments
/// * `start` - Starting number (inclusive)
/// * `end` - Ending number (inclusive)
///
/// # Returns
/// * `Vec<String>` - List of numbers as strings
fn expand_range(start: i32, end: i32) -> Vec<String> {
(start..=end).map(|i| i.to_string()).collect()
}
/// Expand a comma-separated list in the format {item1,item2,item3} to a list of strings
///
/// # Arguments
/// * `list` - Comma-separated string to expand
///
/// # Returns
/// * `Vec<String>` - List of expanded items
fn expand_list(list: &str) -> Vec<String> {
list.split(',').map(|s| s.to_string()).collect()
}
/// Expand a server pattern string with range and list notation into individual hostnames
///
/// Supports two expansion types:
/// - Range expansion: server-[1:5] → server-1, server-2, server-3, server-4, server-5
/// - List expansion: server-{prod,dev} → server-prod, server-dev
///
/// # Arguments
/// * `s` - Pattern string to expand
///
/// # Returns
/// * `Vec<Host>` - List of expanded Host objects
fn expand_string(s: &str) -> Vec<Host> {
fn expand_string(string: String) -> Vec<Host> {
let mut result: Vec<String> = Vec::new();
let mut _result: Vec<String> = Vec::new();
let mut hosts: Vec<Host> = Vec::new();
let mut result = vec![s.to_string()];
// First expand all range expressions [start:end]
while let Some(r) = result.iter().find(|s| s.contains('[')) {
let r = r.clone();
let start = r.find('[').unwrap();
let end = match r[start..].find(']') {
None => {
error!("Error parsing host expression. Wrong range expansion '[a:b]'");
process::exit(1);
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));
}
Some(s) => s + start,
};
let colon = match r[start..end].find(':') {
None => {
error!("Error parsing host expression. Missing colon in range expansion '[a:b]'");
process::exit(1);
}
Some(c) => c + start,
};
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));
}
}
// Then expand all list expressions {item1,item2}
while let Some(r) = result.iter().find(|s| s.contains('{')) {
let r = r.clone();
let start = r.find('{').unwrap();
let end = match r.find('}') {
None => {
error!("Error parsing host expression. Wrong range expansion '{{one,two}}'");
process::exit(1);
}
Some(s) => s,
};
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);
}
}
// Convert all expanded strings to Host objects
for hostname in result {
hosts.push(Host {
name: hostname.to_string(),
@ -273,354 +131,121 @@ fn expand_string(s: &str) -> Vec<Host> {
hosts
}
/// Execute a command on a single host using the system SSH client
///
/// This function runs an SSH command using the system's SSH client,
/// capturing and displaying output in real-time with proper formatting.
///
/// # Arguments
/// * `hostname` - Target server hostname
/// * `username` - SSH username
/// * `command` - Command to execute
/// * `common_suffix` - Optional common suffix for hostname display formatting
/// * `code_only` - Whether to display only exit codes
///
/// # Returns
/// * `Result<i32, String>` - Exit code on success or error message
fn execute_ssh_command(
hostname: &str,
username: &str,
command: &str,
common_suffix: &Option<String>,
code_only: bool,
) -> Result<i32, String> {
let display_name = shorten_hostname(hostname, common_suffix);
// Build the SSH command with appropriate options
let mut ssh_cmd = Command::new("ssh");
ssh_cmd
.arg("-o")
.arg("StrictHostKeyChecking=no")
.arg("-o")
.arg("BatchMode=yes")
.arg(format!("{}@{}", username, hostname))
.arg(command)
.stdout(Stdio::piped())
.stderr(Stdio::piped());
// Execute the command
let mut child = match ssh_cmd.spawn() {
Ok(child) => child,
Err(e) => return Err(format!("Failed to start SSH process: {}", e)),
};
// Function to handle output lines with proper block management
let handle_output = |line: String, display_name: &str, code_only: bool, is_stderr: bool| {
if !code_only {
let mut current_block = CURRENT_BLOCK.lock().unwrap();
// Check if we need to close the previous block and open a new one
match current_block.as_ref() {
Some(open_host) if open_host != display_name => {
// Close the previous block
println!("{}", open_host.yellow());
// Open new block
println!("{}", display_name.yellow());
*current_block = Some(display_name.to_string());
}
None => {
// Open new block
println!("{}", display_name.yellow());
*current_block = Some(display_name.to_string());
}
Some(_) => {
// Same host, continue with current block
}
}
// Print the log line with colored separator based on stream type
let separator = if is_stderr {
"".red()
} else {
"".green()
};
println!(
"{} {} {} {}",
separator,
display_name.yellow(),
separator,
line
);
}
};
// Capture and display stdout in real-time using a dedicated thread
let stdout = child.stdout.take().unwrap();
let display_name_stdout = display_name.clone();
let code_only_stdout = code_only;
let stdout_thread = thread::spawn(move || {
let reader = BufReader::new(stdout);
for line in reader.lines() {
match line {
Ok(line) => {
handle_output(line, &display_name_stdout, code_only_stdout, false);
// false = stdout
}
Err(_) => break,
}
}
});
// Capture and display stderr in real-time using a dedicated thread
let stderr = child.stderr.take().unwrap();
let display_name_stderr = display_name.clone();
let code_only_stderr = code_only;
let stderr_thread = thread::spawn(move || {
let reader = BufReader::new(stderr);
for line in reader.lines() {
match line {
Ok(line) => {
handle_output(line, &display_name_stderr, code_only_stderr, true);
// true = stderr
}
Err(_) => break,
}
}
});
// Wait for command to complete
let status = match child.wait() {
Ok(status) => status,
Err(e) => return Err(format!("Failed to wait for SSH process: {}", e)),
};
// Wait for stdout and stderr threads to complete
stdout_thread.join().unwrap();
stderr_thread.join().unwrap();
// Close the block if this host was the last one to output
if !code_only {
let mut current_block = CURRENT_BLOCK.lock().unwrap();
if let Some(open_host) = current_block.as_ref() {
if open_host == &display_name {
println!("{}", display_name.yellow());
*current_block = None;
}
}
}
// Format exit code with color (green for success, red for failure)
let exit_code = status.code().unwrap_or(-1);
let code_string = if exit_code == 0 {
format!("{}", exit_code.to_string().green())
} else {
format!("{}", exit_code.to_string().red())
};
// For code-only mode, just show hostname and exit code
if code_only {
println!("{}: [{}]", display_name.yellow(), code_string);
}
Ok(exit_code)
}
/// Main entry point for the application
fn main() {
// Initialize logging with minimal formatting (no timestamp, no target)
env_logger::Builder::from_env(Env::default().default_filter_or("info"))
.format_timestamp(None)
.format_target(false)
.init();
// Parse command-line arguments
let args = Args::parse();
// Build the list of target hosts based on user selection method
let hosts = if args.known_hosts {
// Use regex pattern matching against known_hosts file
info!("Using ~/.ssh/known_hosts to build server list.");
let known_hosts = read_known_hosts();
let mut all_hosts = Vec::new();
for expression in args.expression.iter() {
let re = match Regex::new(expression) {
// Build regex
let re = match Regex::new(&args.expression) {
Ok(result) => result,
Err(e) => {
error!("Error parsing regex. {}", e);
process::exit(1);
}
};
let matched: Vec<Host> = known_hosts
.clone()
// match hostnames from known_hosts to regex
known_hosts
.into_iter()
.filter(|r| re.is_match(&r.name.clone()))
.collect();
all_hosts.extend(matched);
}
all_hosts
.collect()
} else {
// Use pattern expansion syntax (ranges and lists)
info!("Using string expansion to build server list.");
let mut all_hosts = Vec::new();
for expression in args.expression.iter() {
all_hosts.extend(expand_string(expression));
}
all_hosts
expand_string(args.expression)
};
// Remove duplicate hosts while preserving original order
// Dedup hosts from known_hosts file
let matched_hosts: Vec<_> = hosts.into_iter().unique().collect();
// Log parallelism setting if not using the default
if args.parallel != 100 {
warn!("Parallelism: {} thread{}", &args.parallel, {
if args.parallel != 1 {
"s."
} else {
"."
}
});
}
// Store hosts with their original indices to preserve ordering
let mut host_with_indices: Vec<(Host, usize)> = Vec::new();
for (idx, host) in matched_hosts.iter().enumerate() {
host_with_indices.push((host.clone(), idx));
}
// Build MasshHostConfig hostnames list
let mut massh_hosts: Vec<MasshHostConfig> = vec![];
let mut hosts_and_ips: HashMap<IpAddr, String> = HashMap::new();
info!("Matched hosts:");
// Perform DNS resolution for all hosts in parallel
// Results are stored with original indices to maintain order
let resolved_ips_with_indices = Arc::new(Mutex::new(Vec::<(String, IpAddr, usize)>::new()));
host_with_indices
.par_iter()
.for_each(|(host, idx)| match lookup_host(&host.name) {
Ok(ips) if !ips.is_empty() => {
let ip = ips[0];
let mut results = resolved_ips_with_indices.lock().unwrap();
results.push((host.name.clone(), ip, *idx));
}
Ok(_) => {
let mut results = resolved_ips_with_indices.lock().unwrap();
results.push((
host.name.clone(),
IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)),
*idx,
));
}
for host in matched_hosts.iter() {
let ip = match lookup_host(&host.name) {
Ok(ip) => ip[0],
Err(_) => {
let mut results = resolved_ips_with_indices.lock().unwrap();
results.push((
host.name.clone(),
IpAddr::V4(std::net::Ipv4Addr::new(0, 0, 0, 0)),
*idx,
));
}
});
// Sort hosts by original index to maintain consistent display order
let mut resolved_hosts = resolved_ips_with_indices.lock().unwrap().clone();
resolved_hosts.sort_by_key(|(_, _, idx)| *idx);
// Display all matched hosts with their resolved IPs
for (hostname, ip, _) in &resolved_hosts {
if ip.is_unspecified() {
error!("DNS resolve failed: {}", hostname.red());
} else {
info!("{} [{}]", hostname, ip);
error!("{} couldn't be resolved.", &host.name.red());
continue;
}
};
info!("{} [{}]", &host.name, ip);
hosts_and_ips.insert(ip, host.name.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,
default_port: 22,
default_user: args.username,
threads: args.parallel as u64,
timeout: 0,
hosts: massh_hosts,
};
let massh = MasshClient::from(&config);
// Filter out hosts that couldn't be resolved
let valid_hosts: Vec<(String, IpAddr, usize)> = resolved_hosts
.into_iter()
.filter(|(_, ip, _)| !ip.is_unspecified())
.collect();
// Exit if no valid hosts remain
if valid_hosts.is_empty() {
error!("No valid hosts to connect to");
process::exit(1);
}
// Find common domain suffix to optimize display
let hostnames: Vec<String> = valid_hosts
.iter()
.map(|(hostname, _, _)| hostname.clone())
.collect();
let common_suffix = find_common_suffix(&hostnames);
// Inform user about display optimization if common suffix found
if let Some(suffix) = &common_suffix {
info!(
"Common domain suffix found: '{}' (will be displayed as '*')",
suffix
);
}
// Ask for confirmation before proceeding (unless --noconfirm is specified)
if !args.noconfirm
&& match Question::new(&*format!(
// Ask for confirmation
if args.noconfirm == true
|| Confirm::new()
.with_prompt(format!(
"Continue on following {} servers?",
&valid_hosts.len()
&config.hosts.len()
))
.confirm()
.interact()
.unwrap()
{
Answer::YES => true,
Answer::NO => {
info!("\n");
info!("Run command on {} servers.", &config.hosts.len());
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() {
let ip: String = host.split('@').collect::<Vec<_>>()[1]
.split(':')
.collect::<Vec<_>>()[0]
.to_string();
let ip = ip.parse::<IpAddr>().unwrap();
info!(
"{}",
hosts_and_ips
.get(&ip)
.unwrap_or(&"Couldn't parse IP".to_string())
.to_string()
.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 {
println!("Code {}", output.exit_status);
} else {
error!("Code {}", output.exit_status);
};
if !args.code {
println!("STDOUT:\n{}", String::from_utf8(output.stdout).unwrap());
println!("STDERR:\n{}", String::from_utf8(output.stderr).unwrap());
}
}
} else {
warn!("Stopped");
process::exit(0);
}
_ => unreachable!(),
}
{
info!("Run command on {} servers.", &valid_hosts.len());
}
// Execute commands using system SSH client
let batch_size = args.parallel as usize;
let mut processed = 0;
while processed < valid_hosts.len() {
let end = std::cmp::min(processed + batch_size, valid_hosts.len());
let batch = &valid_hosts[processed..end];
// Create a thread for each host in the current batch
let mut handles = Vec::new();
for (hostname, _, _) in batch {
let hostname = hostname.clone();
let username = args.username.clone();
let command = args.command.clone();
let common_suffix_clone = common_suffix.clone();
let code_only = args.code;
// Execute SSH command in a separate thread
let handle = thread::spawn(move || {
match execute_ssh_command(
&hostname,
&username,
&command,
&common_suffix_clone,
code_only,
) {
Ok(_) => (),
Err(e) => error!("Error executing command on {}: {}", hostname, e),
}
});
handles.push(handle);
}
// Wait for all threads in this batch to complete
for handle in handles {
handle.join().unwrap();
}
processed = end;
}
}