feat: add ssh.rs
This commit is contained in:
233
ssh-rs/src/main.rs
Normal file
233
ssh-rs/src/main.rs
Normal file
@@ -0,0 +1,233 @@
|
||||
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())
|
||||
}
|
||||
Reference in New Issue
Block a user