From bd1d819e51ff29b5f9dc8b7b4de45032d77eefa9 Mon Sep 17 00:00:00 2001 From: Yao Zongyou Date: Sat, 26 Oct 2024 15:54:17 +0800 Subject: [PATCH] initial version pssh-rs --- Cargo.lock | 491 ++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 16 ++ LICENSE | 21 ++ README.md | 44 ++++ hosts.toml | 18 ++ rust-toolchain.toml | 2 + rustfmt.toml | 6 + src/args.rs | 156 ++++++++++++++ src/main.rs | 223 ++++++++++++++++++++ 9 files changed, 977 insertions(+) create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 LICENSE create mode 100644 README.md create mode 100644 hosts.toml create mode 100644 rust-toolchain.toml create mode 100644 rustfmt.toml create mode 100644 src/args.rs create mode 100644 src/main.rs diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..d58a49b --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,491 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "ansi_term" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d52a9bb7ec0cf484c551830a7ce27bd20d67eac647e1befb56b0be4ee39a55d2" +dependencies = [ + "winapi", +] + +[[package]] +name = "anyhow" +version = "1.0.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c042108f3ed77fd83760a5fd79b53be043192bb3b9dba91d8c574c0ada7850c8" + +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi", + "libc", + "winapi", +] + +[[package]] +name = "autocfg" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "cc" +version = "1.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2e7962b54006dcfcc61cb72735f4d89bb97061dd6a7ed882ec6b8ee53714c6f" +dependencies = [ + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "clap" +version = "2.34.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0610544180c38b88101fecf2dd634b174a62eef6946f84dfc6a7127512b381c" +dependencies = [ + "ansi_term", + "atty", + "bitflags", + "strsim", + "textwrap", + "unicode-width", + "vec_map", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +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.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" + +[[package]] +name = "either" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60b1af1c220855b6ceac025d3f6ecdd2b7c4894bfe9cd9bda4fbb4bc7c0d4cf0" + +[[package]] +name = "heck" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d621efb26863f0e9924c6ac577e8275e5e6b77455db64ffa6c65c904e9e132c" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "hermit-abi" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" +dependencies = [ + "libc", +] + +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.161" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9489c2807c139ffd9c1794f4af0ebe86a828db53ecdc7fea2111d0fed085d1" + +[[package]] +name = "libssh2-sys" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dc8a030b787e2119a731f1951d6a773e2280c660f8ec4b0f5e1505a386e71ee" +dependencies = [ + "cc", + "libc", + "libz-sys", + "openssl-sys", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "libz-sys" +version = "1.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2d16453e800a8cf6dd2fc3eb4bc99b786a9b90c663b8559a5b1a041bf89e472" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "lock_api" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +dependencies = [ + "autocfg", + "scopeguard", +] + +[[package]] +name = "openssl-src" +version = "300.4.0+3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a709e02f2b4aca747929cca5ed248880847c650233cf8b8cdc48f40aaf4898a6" +dependencies = [ + "cc", +] + +[[package]] +name = "openssl-sys" +version = "0.9.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45abf306cbf99debc8195b66b7346498d7b10c210de50418b5ccd7ceba08c741" +dependencies = [ + "cc", + "libc", + "openssl-src", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "parking_lot" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d17b78036a60663b797adeaee46f5c9dfebb86948d1255007a1d6be0271ff99" +dependencies = [ + "instant", + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a2cfe6f0ad2bfc16aefa463b497d5c7a5ecd44a23efa72aa342d90177356dc" +dependencies = [ + "cfg-if", + "instant", + "libc", + "redox_syscall", + "smallvec", + "winapi", +] + +[[package]] +name = "pkg-config" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" + +[[package]] +name = "proc-macro-error" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" +dependencies = [ + "proc-macro-error-attr", + "proc-macro2", + "quote", + "syn 1.0.109", + "version_check", +] + +[[package]] +name = "proc-macro-error-attr" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" +dependencies = [ + "proc-macro2", + "quote", + "version_check", +] + +[[package]] +name = "proc-macro2" +version = "1.0.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f139b0662de085916d1fb67d2b4169d1addddda1919e696f3252b740b629986e" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pssh-rs" +version = "0.5.0" +dependencies = [ + "ansi_term", + "anyhow", + "rayon", + "ssh2", + "structopt", + "toml", +] + +[[package]] +name = "quote" +version = "1.0.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +dependencies = [ + "proc-macro2", +] + +[[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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb5a58c1855b4b6819d59012155603f0b22ad30cad752600aadfcb695265519a" +dependencies = [ + "bitflags", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.213" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ea7893ff5e2466df8d720bb615088341b295f849602c6956047f8f80f0e9bc1" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.213" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e85ad2009c50b58e87caa8cd6dac16bdf511bbfb7af6c33df902396aa480fa5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.85", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "smallvec" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" + +[[package]] +name = "ssh2" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7fe461910559f6d5604c3731d00d2aafc4a83d1665922e280f42f9a168d5455" +dependencies = [ + "bitflags", + "libc", + "libssh2-sys", + "parking_lot", +] + +[[package]] +name = "strsim" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ea5119cdb4c55b55d432abb513a0429384878c15dde60cc77b1c99de1a95a6a" + +[[package]] +name = "structopt" +version = "0.3.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6b5c64445ba8094a6ab0c3cd2ad323e07171012d9c98b0b15651daf1787a10" +dependencies = [ + "clap", + "lazy_static", + "structopt-derive", +] + +[[package]] +name = "structopt-derive" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcb5ae327f9cc13b68763b5749770cb9e048a99bd9dfdfa58d0cf05d5f64afe0" +dependencies = [ + "heck", + "proc-macro-error", + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "2.0.85" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5023162dfcd14ef8f32034d8bcd4cc5ddc61ef7a247c024a33e24e1f24d21b56" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "textwrap" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d326610f408c7a4eb6f51c37c330e496b08506c9457c9d34287ecc38809fb060" +dependencies = [ + "unicode-width", +] + +[[package]] +name = "toml" +version = "0.5.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4f7f0dd8d50a853a531c426359045b1998f04219d88799810762cd4ad314234" +dependencies = [ + "serde", +] + +[[package]] +name = "unicode-ident" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe" + +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[package]] +name = "unicode-width" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dd6e30e90baa6f72411720665d41d89b9a3d039dc45b8faea1ddd07f617f6af" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "vec_map" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1bddf1187be692e79c5ffeab891132dfb0f236ed36a43c7ed39f1165ee20191" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..8b4f694 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "pssh-rs" +version = "0.5.0" +edition = "2021" +description = "pssh-rs is a parallel ssh tool written in rust." +license-file = "LICENSE" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +ssh2 = { version = "0.9", features = ["vendored-openssl"] } +structopt = "0.3" +anyhow = "1.0" +ansi_term = "0.12" +rayon = "1.6" +toml = "0.5" diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..5086051 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2022 Yao Zongyou + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..aa23009 --- /dev/null +++ b/README.md @@ -0,0 +1,44 @@ +pssh-rs is a parallel ssh tool written in rust. + +## Example + +1. Generate config file template: + +```bash +./pssh-rs init +``` +after this command, `hosts.toml` will be generated at the current directory, +change file contents and using --config ./hosts.toml argument to use this file. + +2. run `date` command on default hosts: + +```bash +./pssh-rs --config=./hosts.toml --num_threads=10 run 'date' +``` + +run `date` on all nginx hosts + +```bash +./pssh-rs --config=./hosts.toml --num_threads=10 -s nginx run 'date' +``` + +3. send file to remote hosts: + +```bash +./pssh-rs --config=./hosts.toml send ./hello.txt /tmp/ +``` + +## Install + +just run `cargo install pssh-rs` to install. + +## Building + +pssh-rs can be built with `cargo build --release`, or using the following +command to build statically: + +```bash +sudo apt install musl-tools -y +rustup target add x86_64-unknown-linux-musl +cargo build --target=x86_64-unknown-linux-musl --release +``` diff --git a/hosts.toml b/hosts.toml new file mode 100644 index 0000000..16adcb4 --- /dev/null +++ b/hosts.toml @@ -0,0 +1,18 @@ +username = "root" +password = "123456" +port = 22 +timeout_ms = 10000 +hosts = [ + "192.168.56.101", + "192.168.56.102" +] + +[nginx] +username = "root" +password = "123456" +port = 22 +timeout_ms = 10000 +hosts = [ + "192.168.57.101", + "192.168.57.102" +] diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..2e2b8c8 --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,2 @@ +[toolchain] +channel = "1.82.0" diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..106f2e1 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,6 @@ +# https://rust-lang.github.io/rustfmt/ + +edition = "2021" +use_small_heuristics = "Max" +newline_style = "Unix" +max_width = 120 diff --git a/src/args.rs b/src/args.rs new file mode 100644 index 0000000..2513dab --- /dev/null +++ b/src/args.rs @@ -0,0 +1,156 @@ +use anyhow::{anyhow, Context}; +use std::path::PathBuf; +use structopt::StructOpt; +use toml::Value; + +const DEFAULT_SSH_PORT: u16 = 22; +const DEFAULT_SSH_USERNAME: &str = "root"; +const DEFAULT_SSH_TIMEOUT_MS: u32 = 3000; + +#[derive(Clone, Debug, StructOpt)] +#[structopt(name = "pssh-rs", about = "pssh-rs is a parallel ssh tool written in rust")] +pub struct CommandLineArgs { + /// toml file for config + #[structopt(parse(from_os_str), short, long, default_value = "./hosts.toml")] + config: PathBuf, + + /// section in toml file + #[structopt(short = "s", long)] + sections: Option>, + + #[structopt(subcommand)] + pub command: Command, + + /// The number of threads. + #[structopt(short, long = "num_threads", default_value = "1")] + pub num_threads: usize, + + /// Keep the output stable order with designated hosts. + #[structopt(short = "k", long = "keep_stable")] + pub keep_stable: bool, +} + +#[derive(Clone, Debug, StructOpt)] +pub enum Command { + /// Init local hosts.toml config file. + Init, + + /// Run commands on the remote hosts. + Run { + /// The command to run remotely + command: String, + }, + + /// Send file to the remote hosts. + Send { + /// local source file path to send + #[structopt(parse(from_os_str))] + source_fpath: PathBuf, + + /// destination file path + #[structopt(parse(from_os_str))] + target_fpath: PathBuf, + }, +} + +#[derive(Clone, Debug)] +pub struct HostInfo { + pub host: String, + pub username: String, + pub password: String, + pub port: u16, + pub timeout_ms: u32, +} + +impl CommandLineArgs { + pub fn get_hosts(&self) -> anyhow::Result> { + let str = std::fs::read_to_string(&self.config)?; + let value = str.parse::()?; + + let Value::Table(table) = value else { + return Err(anyhow!("illegal toml format: content of toml should be a table")); + }; + + let Some(ref sections) = self.sections else { + return get_hosts_from_table(&table); + }; + + if sections.is_empty() { + return get_hosts_from_table(&table); + } + + let mut res = vec![]; + + for section in sections { + let Some(section_value) = table.get(section) else { + return Err(anyhow!("no {} section in the toml file", section)); + }; + + let Value::Table(section_table) = section_value else { + return Err(anyhow!("illegal section format: content of section should be a table: {}", section)); + }; + + let mut hosts = get_hosts_from_table(section_table)?; + res.append(&mut hosts); + } + + Ok(res) + } +} + +fn get_hosts_from_table(table: &toml::value::Table) -> anyhow::Result> { + let mut res = vec![]; + + let username = get_username(table.get("username"))?.unwrap_or_else(|| DEFAULT_SSH_USERNAME.to_string()); + let password = get_password(table.get("password"))?.unwrap_or_default(); + let port = get_port(table.get("port"))?.unwrap_or(DEFAULT_SSH_PORT); + let timeout_ms = get_timeout_ms(table.get("timeout_ms"))?.unwrap_or(DEFAULT_SSH_TIMEOUT_MS); + + for host in table.get("hosts").iter().flat_map(|a| a.as_array()).flatten().flat_map(|v| v.as_str()) { + res.push(HostInfo { + username: username.clone(), + password: password.clone(), + port, + host: host.to_string(), + timeout_ms, + }) + } + + Ok(res) +} + +fn get_username(value: Option<&Value>) -> anyhow::Result> { + let Some(value) = value else { + return Ok(None); + }; + + let value = value.as_str().ok_or_else(|| anyhow!("username should be a string"))?; + Ok(Some(value.to_string())) +} + +fn get_password(value: Option<&Value>) -> anyhow::Result> { + let Some(value) = value else { + return Ok(None); + }; + + let value = value.as_str().ok_or_else(|| anyhow!("password should be a string"))?; + Ok(Some(value.to_string())) +} + +fn get_port(value: Option<&Value>) -> anyhow::Result> { + let Some(value) = value else { + return Ok(None); + }; + + let value = value.as_integer().ok_or_else(|| anyhow!("port should be an u16"))?; + Ok(Some(value.try_into().context("port should be in the range [0, 65535]")?)) +} + +fn get_timeout_ms(value: Option<&Value>) -> anyhow::Result> { + let Some(value) = value else { + return Ok(None); + }; + + let value = value.as_integer().ok_or_else(|| anyhow!("timeout_ms should be an u32"))?; + Ok(Some(value.try_into().context("timeout_ms should be valid u32")?)) +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..9ac4e30 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,223 @@ +use ansi_term::Colour::{Green, Red}; +use anyhow::Context; +use args::HostInfo; +use rayon::prelude::*; +use ssh2::Session; +use std::fs::metadata; +use std::fs::File; +use std::io::prelude::*; +use std::io::Write; +use std::net::TcpStream; +use std::os::unix::prelude::PermissionsExt; +use std::path::Path; +use std::sync::mpsc::sync_channel; +use std::time::Duration; +use structopt::StructOpt; + +mod args; + +const DEFAULT_CONFIG_FPATH: &str = "./hosts.toml"; +const DEFAULT_CONFIG_TMPL: &str = r#"username = "root" +password = "123456" +port = 22 +timeout_ms = 10000 +hosts = [ + "192.168.56.101", + "192.168.56.102" +] + +[nginx] +username = "root" +password = "123456" +port = 22 +timeout_ms = 10000 +hosts = [ + "192.168.57.101", + "192.168.57.102" +] +"#; + +fn main() -> anyhow::Result<()> { + let args = args::CommandLineArgs::from_args(); + + if let args::Command::Init = args.command { + if std::fs::exists(DEFAULT_CONFIG_FPATH)? { + println!("Save the following contents to a file (default to {}) and then", DEFAULT_CONFIG_FPATH); + println!("using --config {} to use this config file.", DEFAULT_CONFIG_FPATH); + println!(); + println!("{}", DEFAULT_CONFIG_TMPL); + } else { + std::fs::write(DEFAULT_CONFIG_FPATH, DEFAULT_CONFIG_TMPL)?; + println!("default config template has been written to {}", DEFAULT_CONFIG_FPATH); + println!("modify this file and using --config {} to use this config file", DEFAULT_CONFIG_FPATH) + } + return Ok(()); + } + + let hosts = args.get_hosts()?; + let (sender, receiver) = sync_channel(hosts.len()); + + std::thread::scope(|s| { + s.spawn({ + let hosts = hosts.clone(); + move || { + print_main(args.keep_stable, &hosts, receiver); + } + }); + + rayon::ThreadPoolBuilder::new().num_threads(args.num_threads).build_global().unwrap(); + hosts.par_iter().enumerate().for_each(|(index, host)| { + let result = match &args.command { + args::Command::Run { command } => run_command(host, command), + args::Command::Send { source_fpath, target_fpath } => send_file(host, source_fpath, target_fpath), + _ => { + panic!("shoud not come here") + } + }; + + sender.send((index, host, result)).unwrap(); + }); + + drop(sender); + }); + + Ok(()) +} + +enum Outcome { + RunCommandOutcome(RunCommandOutcome), + SendFileOutcome, +} + +struct RunCommandOutcome { + exit_status: i32, + out: Vec, + err: Vec, +} + +fn run_command(host: &HostInfo, command: &str) -> anyhow::Result { + let addr = format!("{}:{}", host.host, host.port); + let tcp = TcpStream::connect_timeout(&addr.parse()?, Duration::from_millis(host.timeout_ms as u64))?; + let mut sess = Session::new()?; + sess.set_tcp_stream(tcp); + sess.set_timeout(host.timeout_ms); + sess.handshake()?; + sess.set_timeout(0); + sess.userauth_password(&host.username, &host.password)?; + + let mut channel = sess.channel_session()?; + let _ = channel.setenv("PSSH_RS_IP", &host.host); + let _ = channel.setenv("PSSH_RS_PORT", &host.port.to_string()); + channel.exec(command)?; + + let (mut out, mut err) = (vec![], vec![]); + + std::thread::scope(|s| { + s.spawn(|| { + if let Err(err) = channel.stream(0).take(1024 * 1024).read_to_end(&mut out) { + println!("failed to read: {err}"); + } + }); + s.spawn(|| { + if let Err(err) = channel.stderr().take(1024 * 1024).read_to_end(&mut err) { + println!("failed to read: {err}"); + } + }); + }); + + channel.wait_close()?; + + let exit_status = channel.exit_status()?; + Ok(Outcome::RunCommandOutcome(RunCommandOutcome { exit_status, out, err })) +} + +fn send_file(host: &HostInfo, source_fpath: &Path, target_fpath: &Path) -> anyhow::Result { + let addr = format!("{}:{}", host.host, host.port); + let tcp = TcpStream::connect_timeout(&addr.parse()?, Duration::from_millis(host.timeout_ms as u64))?; + let mut sess = Session::new()?; + sess.set_tcp_stream(tcp); + sess.set_timeout(host.timeout_ms); + sess.handshake()?; + sess.set_timeout(0); + sess.userauth_password(&host.username, &host.password)?; + + let attr = metadata(source_fpath).with_context(|| format!("stat local {}", source_fpath.to_string_lossy()))?; + let mode = attr.permissions().mode() & 0o777; + + let mut remote_file = sess.scp_send(target_fpath, mode as i32, attr.len(), None)?; + + let mut file = File::open(source_fpath)?; + std::io::copy(&mut file, &mut remote_file)?; + + // Close the channel and wait for the whole content to be transferred + remote_file.send_eof()?; + remote_file.wait_eof()?; + remote_file.close()?; + remote_file.wait_close()?; + + Ok(Outcome::SendFileOutcome) +} + +fn print_main( + keep_stable: bool, + hosts: &[args::HostInfo], + receiver: std::sync::mpsc::Receiver<(usize, &args::HostInfo, anyhow::Result)>, +) { + let mut results: Vec<_> = std::iter::repeat_with(|| None).take(hosts.len()).collect(); + let mut print_index: usize = 0; + + loop { + let Ok((index, host, result)) = receiver.recv() else { + break; + }; + + if !keep_stable { + print_outcome(host, &result).unwrap(); + continue; + } + + results[index] = Some((host, result)); + while print_index < results.len() { + match results[print_index] { + Some((host, ref result)) => { + print_outcome(host, result).unwrap(); + print_index += 1; + } + None => { + break; + } + } + } + } +} + +fn print_outcome(host: &HostInfo, result: &anyhow::Result) -> anyhow::Result<()> { + let addr = format!("{}:{}", host.host, host.port); + + match result { + Ok(Outcome::RunCommandOutcome(command_outcome)) => { + print_command_outcome(&addr, command_outcome)?; + } + Ok(Outcome::SendFileOutcome) => { + println!("{}", Green.paint(format!("[{addr} OK]"))); + } + Err(err) => { + println!("{}", Red.paint(format!("[{addr} ERROR: {err:#}]"))); + } + } + + Ok(()) +} + +fn print_command_outcome(addr: &str, outcomd: &RunCommandOutcome) -> anyhow::Result<()> { + if outcomd.exit_status == 0 { + println!("{}", Green.paint(format!("[{addr} OK]"))); + } else { + println!("{}", Red.paint(format!("[{addr} ERROR: exit with {}]", outcomd.exit_status))); + } + + std::io::stdout().write_all(&outcomd.out)?; + std::io::stdout().write_all(&outcomd.err)?; + + Ok(()) +}