clean up code to make it easier for more services and add clap

Cargo.lock

@@ -2,6 +2,56 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "anstream"
+version = "0.6.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "301af1932e46185686725e0fad2f8f2aa7da69dd70bf6ecc44d6b703844a3933"
+dependencies = [
+ "anstyle",
+ "anstyle-parse",
+ "anstyle-query",
+ "anstyle-wincon",
+ "colorchoice",
+ "is_terminal_polyfill",
+ "utf8parse",
+]
+
+[[package]]
+name = "anstyle"
+version = "1.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "862ed96ca487e809f1c8e5a8447f6ee2cf102f846893800b20cebdf541fc6bbd"
+
+[[package]]
+name = "anstyle-parse"
+version = "0.2.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2"
+dependencies = [
+ "utf8parse",
+]
+
+[[package]]
+name = "anstyle-query"
+version = "1.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6c8bdeb6047d8983be085bab0ba1472e6dc604e7041dbf6fcd5e71523014fae9"
+dependencies = [
+ "windows-sys",
+]
+
+[[package]]
+name = "anstyle-wincon"
+version = "3.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "403f75924867bb1033c59fbf0797484329750cfbe3c4325cd33127941fabc882"
+dependencies = [
+ "anstyle",
+ "once_cell_polyfill",
+ "windows-sys",
+]
+
 [[package]]
 name = "autocfg"
 version = "1.5.0"
@@ -29,12 +79,70 @@ version = "1.0.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268"
 
+[[package]]
+name = "clap"
+version = "4.5.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "40b6887a1d8685cebccf115538db5c0efe625ccac9696ad45c409d96566e910f"
+dependencies = [
+ "clap_builder",
+ "clap_derive",
+]
+
+[[package]]
+name = "clap_builder"
+version = "4.5.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e0c66c08ce9f0c698cbce5c0279d0bb6ac936d8674174fe48f736533b964f59e"
+dependencies = [
+ "anstream",
+ "anstyle",
+ "clap_lex",
+ "strsim",
+]
+
+[[package]]
+name = "clap_derive"
+version = "4.5.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d2c7947ae4cc3d851207c1adb5b5e260ff0cca11446b1d6d1423788e442257ce"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "clap_lex"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675"
+
+[[package]]
+name = "colorchoice"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75"
+
 [[package]]
 name = "emojis-rs"
 version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e5c5d55d32597632ca64765dd57c1e71aaa51e69f663671fdc955950062c38ff"
 
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "is_terminal_polyfill"
+version = "1.70.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf"
+
 [[package]]
 name = "libc"
 version = "0.2.174"
@@ -77,6 +185,12 @@ dependencies = [
  "scopeguard",
 ]
 
+[[package]]
+name = "once_cell_polyfill"
+version = "1.70.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad"
+
 [[package]]
 name = "openssl-sys"
 version = "0.9.109"
@@ -118,6 +232,24 @@ version = "0.3.32"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c"
 
+[[package]]
+name = "proc-macro2"
+version = "1.0.95"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d"
+dependencies = [
+ "proc-macro2",
+]
+
 [[package]]
 name = "redox_syscall"
 version = "0.5.13"
@@ -149,6 +281,7 @@ checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
 name = "ssh-test"
 version = "0.1.0"
 dependencies = [
+ "clap",
  "emojis-rs",
  "ssh2",
 ]
@@ -165,12 +298,50 @@ dependencies = [
  "parking_lot",
 ]
 
+[[package]]
+name = "strsim"
+version = "0.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f"
+
+[[package]]
+name = "syn"
+version = "2.0.104"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "17b6f705963418cdb9927482fa304bc562ece2fdd4f616084c50b7023b435a40"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
+
+[[package]]
+name = "utf8parse"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
+
 [[package]]
 name = "vcpkg"
 version = "0.2.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
 
+[[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets",
+]
+
 [[package]]
 name = "windows-targets"
 version = "0.52.6"

Cargo.toml

@@ -4,5 +4,6 @@ version = "0.1.0"
 edition = "2024"
 
 [dependencies]
+clap = { version = "4.5.40", features = ["derive"] }
 emojis-rs = "0.1.3"
 ssh2 = "0.9.5"

src/main.rs

@@ -1,54 +1,58 @@
 use std::io::prelude::*;
 use std::net::TcpStream;
 use ssh2::Session;
+use clap::Parser;
 use emojis_rs::*;
 
-const USER: &str = "aaronh";
+//const USER: &str = "aaronh";
-const EDI: &str = "100.94.173.5:22";
+//const EDI: &str = "100.94.173.5:22";
+
+#[derive(Parser, Debug)]
+#[command(version, about, long_about = None)]
+struct Args {
+    #[arg(short, long)]
+    username: String,
+    host: String,
+}
 
 fn main() {
+    let args = Args::parse();
     
-    // Connect to the local SSH server
+    let tcp = TcpStream::connect(args.host).expect("Failed to connect to SSH server");
-    let tcp = TcpStream::connect(EDI).unwrap();
+    let mut sess = Session::new().expect("Failed to create SSH session");
-    let mut sess = Session::new().unwrap();
     sess.set_tcp_stream(tcp);
-    sess.handshake().unwrap();
+    sess.handshake().expect("SSH handshake failed");
     
-    // Try to authenticate with the first identity in the agent.
+    sess.userauth_agent(&args.username).expect("SSH authentication failed");
-    sess.userauth_agent(USER).unwrap();
 
     if sess.authenticated() {
         println!("Authentication successful! {EMOJI_CHECK}");
     } else {
         println!("Authentication failed! {EMOJI_CROSS}");
+        return;
     }
 
     println!("-------------------------");
-    let mut gollum_check = sess.channel_session().unwrap();
-    gollum_check.exec("systemctl is-active gollum").unwrap();
-    let mut s = String::new();
-    gollum_check.read_to_string(&mut s).unwrap();
-    gollum_check.wait_close().unwrap();
 
-    let exit_code = gollum_check.exit_status().unwrap();
+    for service in ["gollum", "wastebin", "hydra-server"] {
+        check_service_status(&sess, service);
+    }
+}
+
+fn check_service_status(sess: &Session, service: &str) {
+    let mut channel = sess.channel_session().expect("Failed to open SSH channel");
+    channel.exec(&format!("systemctl is-active {service}")).expect("Failed to run systemctl");
+    
+    let mut output = String::new();
+    channel.read_to_string(&mut output).expect("Failed to read channel output");
+    channel.wait_close().expect("Failed to close SSH channel");
+
+    let exit_code = channel.exit_status().expect("Failed to get exit status");
 
     if exit_code == 0 {
-        println!("Gollum service is active {EMOJI_CHECK}");
+        println!("{service} service is active {EMOJI_CHECK}");
     } else {
-        println!("Gollum service is inactive {EMOJI_CROSS}(exit code: {})", exit_code);
+        println!("{service} service is inactive {EMOJI_CROSS} (exit code: {exit_code})");
-    }
-
-    let mut wastebin_check = sess.channel_session().unwrap();
-    wastebin_check.exec("systemctl is-active wastebin").unwrap();
-    let mut s = String::new();
-    wastebin_check.read_to_string(&mut s).unwrap();
-    wastebin_check.wait_close().unwrap();
-
-    let exit_code = wastebin_check.exit_status().unwrap();
-
-    if exit_code == 0 {
-        println!("Wastebin service is active {EMOJI_CHECK}");
-    } else {
-        println!("Wastebin service is inactive {EMOJI_CROSS}(exit code: {})", exit_code);
     }
 }
+