moshy-rust/src/main.rs

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 <kiniou@gmail.com>")
        .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("<hostname> 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<String> = 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<Vec<String>> = None;
    if let Some(cmd_args) = matches.values_of("cmd") {
        let cmd_vec: Vec<String> = 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<String>,
    cmd: Option<Vec<String>>,
}

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
}