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

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 EDI: &str = "100.94.173.5:22";
+//const USER: &str = "aaronh";
+//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() {
-
-    // Connect to the local SSH server
-    let tcp = TcpStream::connect(EDI).unwrap();
-    let mut sess = Session::new().unwrap();
-    sess.set_tcp_stream(tcp);
-    sess.handshake().unwrap();
+    let args = Args::parse();
     
-    // Try to authenticate with the first identity in the agent.
-    sess.userauth_agent(USER).unwrap();
+    let tcp = TcpStream::connect(args.host).expect("Failed to connect to SSH server");
+    let mut sess = Session::new().expect("Failed to create SSH session");
+    sess.set_tcp_stream(tcp);
+    sess.handshake().expect("SSH handshake failed");
+    
+    sess.userauth_agent(&args.username).expect("SSH authentication failed");
 
     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();
-    
-    if exit_code == 0 {
-        println!("Gollum service is active {EMOJI_CHECK}");
-    } else {
-        println!("Gollum 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);
+    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!("{service} service is active {EMOJI_CHECK}");
+    } else {
+        println!("{service} service is inactive {EMOJI_CROSS} (exit code: {exit_code})");
+    }
+}
+