scriptbase/ssh-rs/src/main.rs

use clap::Parser;
use rust_util::{
    debugging, failure_and_exit, information, opt_result, simple_error, success, util_cmd, util_env,
    util_file, util_term, XResult,
};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::process::Command;

#[derive(Debug, Parser)]
#[command(name = "ssh-rs", bin_name = "ssh.rs")]
#[command(about = "SSH util", long_about = None)]
struct SshRsArgs {
    /// Forward agent
    #[arg(long, short = 'f')]
    pub forward_agent: Option<bool>,
    /// Proxy
    #[arg(long, short = 'p')]
    pub proxy: Option<bool>,
    /// [username@]host
    pub username_and_host: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
struct SshConfig {
    pub default_forward_agent: Option<bool>,
    pub default_proxy: Option<bool>,
    pub default_username: Option<String>,
    pub profiles: BTreeMap<String, SshProfile>,
}

#[derive(Debug, Serialize, Deserialize)]
struct SshProfile {
    pub default_username: Option<String>,
    pub alias: Option<Vec<String>>,
    pub host: String,
    pub proxy: Option<bool>,
    pub forward_agent: Option<bool>,
    pub comment: Option<String>,
}

impl SshConfig {
    fn find_profiles(&self, profile: &str) -> Vec<(&String, &SshProfile)> {
        let mut found = vec![];
        for (k, v) in &self.profiles {
            if k == profile {
                found.push((k, v));
            } else {
                if let Some(alias) = &v.alias {
                    if alias.contains(&profile.to_string()) {
                        found.push((k, v));
                    }
                }
            }
        }
        found
    }
}

const ENV_SSH_RS_CONFIG_FILE: &str = "SSH_RS_CONFIG_FILE";
const SSH_RS_CONFIG_FILE: &str = "~/.config/ssh-rs-config.json";
const CMD_SSH: &str = "ssh";

fn main() -> XResult<()> {
    let args = SshRsArgs::parse();
    let ssh_rs_config = load_ssh_rs_config()?;

    let username_and_host = match args.username_and_host {
        None => {
            success!("Total {} server(s):", ssh_rs_config.profiles.len());
            let mut max_profile_id_len = 0_usize;
            let mut max_host_len = 0_usize;
            for (profile_id, profile) in &ssh_rs_config.profiles {
                if profile_id.len() > max_profile_id_len {
                    max_profile_id_len = profile_id.len();
                }
                if profile.host.len() > max_host_len {
                    max_host_len = profile.host.len();
                }
            }
            for (profile_id, profile) in &ssh_rs_config.profiles {
                let mut features = vec![];
                if let Some(true) = profile.forward_agent {
                    features.push("forward_agent");
                }
                if let Some(true) = profile.proxy {
                    features.push("proxy");
                }
                println!(
                    "- {} : {}{} {}{}{} # {}{}",
                    pad(profile_id, max_profile_id_len),
                    util_term::GREEN,
                    pad(&profile.host, max_host_len),
                    util_term::YELLOW,
                    match &profile.alias {
                        None => {
                            "".to_string()
                        }
                        Some(alias) => {
                            format!("alias: [{}]", alias.join(", "))
                        }
                    },
                    util_term::END,
                    profile.comment.clone().unwrap_or_else(|| "-".to_string()),
                    if features.is_empty() {
                        "".to_string()
                    } else {
                        format!(" ;[{}]", features.join(", "))
                    }
                );
            }

            let mut features = vec![];
            if let Some(true) = ssh_rs_config.default_forward_agent {
                features.push("forward_agent");
            }
            if let Some(true) = ssh_rs_config.default_proxy {
                features.push("proxy");
            }
            if !features.is_empty() {
                println!();
                information!("Global default features: [{}]", features.join(", "));
            }

            return Ok(());
        }
        Some(username_and_host) => username_and_host,
    };

    let (username, host) = parse_username_and_host(&username_and_host)?;

    let profiles = ssh_rs_config.find_profiles(&host);
    if profiles.is_empty() {
        return simple_error!("Profile not found");
    }
    if profiles.len() > 1 {
        let profile_names = profiles
            .iter()
            .map(|p| p.0.to_string())
            .collect::<Vec<_>>()
            .join(", ");
        return simple_error!("Multiple profiles found: {}", profile_names);
    }
    let profile = profiles[0].1;

    debugging!("Found profile: {:#?}", profile);

    let ssh_forward_agent = args.forward_agent.clone().unwrap_or_else(|| {
        profile.forward_agent.unwrap_or_else(|| //-
                ssh_rs_config.default_forward_agent.unwrap_or(false))
    });
    let ssh_proxy = args.proxy.clone().unwrap_or_else(|| {
        profile.proxy.unwrap_or_else(|| //- 
                ssh_rs_config.default_proxy.unwrap_or(false))
    });
    let ssh_username = username.unwrap_or_else(|| {
        profile.default_username.clone().unwrap_or_else(|| {
            ssh_rs_config
                .default_username
                .clone()
                .unwrap_or_else(|| "root".to_string())
        })
    });
    let ssh_host = profile.host.clone();

    let mut cmd = Command::new(CMD_SSH);
    if ssh_forward_agent {
        cmd.args(&["-o", "ForwardAgent=yes"]);
    }
    if ssh_proxy {
        cmd.args(&["-o", "ProxyCommand=nc -X 5 -x 127.0.0.1:1080 %h %p"]);
    }
    cmd.arg(&format!("{}@{}", ssh_username, ssh_host));

    success!(
        "{} {}",
        CMD_SSH,
        cmd.get_args()
            .map(|s| s.to_string_lossy().to_string())
            .collect::<Vec<_>>()
            .join(" ")
    );
    match util_cmd::run_command_and_wait(&mut cmd) {
        Ok(exit_status) => {
            debugging!("Command exited with status: {}", exit_status);
            if !exit_status.success() {
                failure_and_exit!("Exit with error: {}", exit_status);
            }
        }
        Err(err) => {
            failure_and_exit!("SSH command failed: {}", err);
        }
    }
    Ok(())
}

fn pad(str: &str, width: usize) -> String {
    if str.len() >= width {
        return str.to_string();
    }
    format!("{}{}", str, " ".repeat(width - str.len()))
}

fn parse_username_and_host(username_and_host: &str) -> XResult<(Option<String>, String)> {
    if username_and_host.is_empty() {
        return simple_error!("Empty username@host");
    }
    let username_and_host_parts = username_and_host.split("@").collect::<Vec<_>>();
    if username_and_host_parts.len() == 1 {
        return Ok((None, username_and_host_parts[0].to_string()));
    }
    if username_and_host_parts.len() > 2 {
        return simple_error!("Bad username@host: {}", username_and_host);
    }
    Ok((
        Some(username_and_host_parts[0].to_string()),
        username_and_host_parts[1].to_string(),
    ))
}

fn load_ssh_rs_config() -> XResult<SshConfig> {
    let config_file = get_ssh_rs_config_file();
    let config_content = util_file::read_file_content(&config_file)?;
    let config: SshConfig = opt_result!(
        serde_json::from_str(config_content.as_str()),
        "Parse config failed: {}"
    );
    Ok(config)
}

fn get_ssh_rs_config_file() -> String {
    util_env::env_var(ENV_SSH_RS_CONFIG_FILE).unwrap_or_else(|| SSH_RS_CONFIG_FILE.to_string())
}