1

Signed-off-by: Dmytro Stanchiev <git@dmytros.dev>

src/main.rs

@@ -1,3 +1,220 @@
+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() {
-    println!("Hello, world!");
+    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
 }