use std::collections::HashMap; use std::env; use std::path::Path; use std::process::{self, Command}; use std::str; use clap::{App, Arg, ArgMatches}; use config::{Config, File, FileFormat}; use log::{LevelFilter, debug, info}; use rand::Rng; use shell_words; fn main() { let matches = App::new("moshy") .version("1.0") .author("Kevin Roy ") .about("Mosh powered with profiles") .arg( Arg::with_name("debug") .short('d') .long("debug") .help("Debug moshy to see what's done in background"), ) .arg( Arg::with_name("dry-run") .short('n') .long("dry-run") .help("Shows what will be executed"), ) .subcommand( App::new("list") .about("List every remote hostname configured as profile") .arg( Arg::with_name("verbose") .short('v') .long("verbose") .help("Increase verbosity"), ), ) .arg( Arg::with_name("hostname") .help(" parameter match configuration sections and can be formatted like the following: hostname, user@hostname, hostname:flavor, user@hostname:flavor") .required(false), ) .arg( Arg::with_name("cmd") .multiple(true) .help("Extra arguments which will be executed to the remote mosh session") .required(false), ) .get_matches(); // Set up logging let mut log_level = LevelFilter::Info; if matches.is_present("debug") { log_level = LevelFilter::Debug; } env_logger::Builder::new().filter_level(log_level).init(); debug!("Got following arguments:\n{:?}", matches); // Config paths let config_paths = vec![ "/etc/moshy.conf".to_string(), format!( "{}/.moshy.conf", env::var("HOME").unwrap_or_else(|_| "".to_string()) ), ]; // Read config let mut config_builder = Config::builder(); for path in &config_paths { if Path::new(path).exists() { config_builder = config_builder.add_source(File::with_name(path).format(FileFormat::Ini)); } } let config = config_builder.build().unwrap_or_else(|e| { eprintln!("Error reading config: {}", e); process::exit(1); }); if let Some(list_matches) = matches.subcommand_matches("list") { let sections: Vec = config .get_table("") .unwrap_or_default() .keys() .cloned() .collect(); if list_matches.is_present("verbose") { // In verbose mode, perhaps print more details, but original only lists sections println!("{}", sections.join("\n")); } else { println!("{}", sections.join("\n")); } process::exit(0); } let hostname_arg = matches.value_of("hostname").unwrap_or(""); let mut host_setup = parse_hostname(hostname_arg, &config); debug!("Required host setup:\n{:?}", host_setup); let mut cmd = vec!["/usr/bin/env".to_string(), "mosh".to_string()]; if let Some(ports) = &host_setup.ports { let ports_split: Vec<&str> = ports.split(':').collect(); let port = if ports_split.len() > 1 { let start: u16 = ports_split[0].parse().unwrap_or(0); let end: u16 = ports_split[1].parse().unwrap_or(0); rand::thread_rng().gen_range(start..=end) } else { ports_split[0].parse().unwrap_or(0) }; cmd.push(format!("--port={}", port)); } cmd.push(host_setup.host.clone()); let mut extra_cmd: Option> = None; if let Some(cmd_args) = matches.values_of("cmd") { let cmd_vec: Vec = cmd_args.map(|s| s.to_string()).collect(); if cmd_vec.len() == 1 { extra_cmd = shell_words::split(&cmd_vec[0]).ok(); } else { extra_cmd = Some(cmd_vec); } } if let Some(ref mut cmd_setup) = host_setup.cmd { if let Some(extra) = extra_cmd { cmd_setup.extend(extra); } } else if let Some(extra) = extra_cmd { host_setup.cmd = Some(extra); } if let Some(cmd_setup) = &host_setup.cmd { cmd.push("--".to_string()); cmd.extend(cmd_setup.clone()); } let log_msg = format!("Will execute the following command:\n {}", cmd.join(" ")); if matches.is_present("dry-run") { info!("{}", log_msg); println!("{:?}", cmd); } else { debug!("{}", log_msg); let mut child = Command::new(&cmd[0]) .args(&cmd[1..]) .spawn() .expect("Failed to execute command"); child.wait().expect("Command wasn't running"); } process::exit(0); } #[derive(Debug, Clone)] struct HostSetup { host: String, ports: Option, cmd: Option>, } fn parse_hostname(hostname_argument: &str, config: &Config) -> HostSetup { let mut host_setup = HostSetup { host: "".to_string(), ports: None, cmd: None, }; if !hostname_argument.is_empty() { debug!("Parsing hostname argument '{}'", hostname_argument); let connection: Vec<&str> = hostname_argument.split('@').collect(); let hostname = if connection.len() > 1 { connection[1] } else { connection[0] }; let mut profiles = Vec::new(); if config.get_table(hostname).is_ok() { profiles.push(hostname.to_string()); } let flavor: Vec<&str> = hostname.split(':').collect(); if flavor.len() > 1 { let hostname_no_flavor = flavor[0]; if config.get_table(hostname_no_flavor).is_ok() { profiles.push(hostname_no_flavor.to_string()); } } debug!("Found hostname profile(s) {}", hostname); debug!("Profile(s) to look for values :\n{:?}", profiles); host_setup.host = hostname.to_string(); profiles.reverse(); for profile in profiles { debug!("Get values from {}", profile); if let Ok(table) = config.get_table(&profile) { debug!("profile config:\n{:?}", table); if let Some(ports) = table.get("port").and_then(|v| v.clone().into_string().ok()) { host_setup.ports = Some(ports); } debug!("ports:{:?}", host_setup.ports); if let Some(cmd_str) = table.get("cmd").and_then(|v| v.clone().into_string().ok()) { if let Ok(cmd_vec) = shell_words::split(&cmd_str) { host_setup.cmd = Some(cmd_vec); } } debug!("command:{:?}", host_setup.cmd); } } } host_setup }