1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

use std::io::Result;
use std::path::{Path, PathBuf};
use std::process::{Output, Child, ExitStatus, Command};
use shlex::Shlex;

/// This describes a remote machine that we can connect to.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Remote {
    pub addr: String,
    pub port: u16,
}

/// This is solely implemented for the `Command` type.
/// It adds three methods which allow the "command" type
/// to behave the same on a remote machine as it does on a local.
pub trait RemoteCommand {
    /// Acts identically to the `output` method on `Command`,
    /// except that it runs the command on the remote machine.
    fn remote_output(&mut self, remote: &Remote) -> Result<Output>;

    /// Acts identically to the `spawn` method on `Command`,
    /// except that it runs the command on the remote machine.
    fn remote_spawn(&mut self, remote: &Remote) -> Result<Child>;

    /// Acts identically to the `status` method on `Command`,
    /// except that it runs the command on the remote machine.
    fn remote_status(&mut self, remote: &Remote) -> Result<ExitStatus>;
}

fn reconstruct_ssh_command(remote: &Remote, command: &Command) -> Command {
    let mut cmd = Command::new("ssh");
    cmd.arg("-p").arg(remote.port.to_string());
    cmd.arg(remote.addr.clone());

    // Check whether ~/.ssh/cs6991/cs6991-id exists
    let path: PathBuf = [home::home_dir().unwrap().to_str().unwrap(), ".ssh", "cs6991", "cs6991-id"].iter().collect();
    if path.is_file() {
        cmd.arg("-i");
        cmd.arg(path.as_os_str());
    }

    cmd.arg("--");
    cmd.arg(command.get_program());
    cmd.args(command.get_args());

    cmd
}

impl RemoteCommand for Command {
    fn remote_output(&mut self, remote: &Remote) -> Result<Output> {
        let mut cmd = reconstruct_ssh_command(remote, self);
        cmd.output()
    }

    fn remote_spawn(&mut self, remote: &Remote) -> Result<Child> {
        let mut cmd = reconstruct_ssh_command(remote, self);
        let mut cmd = cmd.stdin(std::process::Stdio::piped());
        let mut cmd = cmd.stdout(std::process::Stdio::piped());
        cmd.spawn()
    }

    fn remote_status(&mut self, remote: &Remote) -> Result<ExitStatus> {
        let mut cmd = reconstruct_ssh_command(remote, self);
        cmd.status()
    }
}

/// This takes a line of input, and splits it into a vector of commands.
pub fn parse_line(s: &str) -> Option<Vec<Vec<String>>> {
    let mut cmd = vec![];
    let mut cmds = vec![];
    for token in shlex::split(s)? {
        if token == ";" && !cmd.is_empty() {
            cmds.push(cmd);
            cmd = vec![];
        } else {
            let starts_with_split = token.starts_with(";");
            let ends_with_split = token.ends_with(";");
            let token = token.trim_matches(';').to_string();
            if starts_with_split && !cmd.is_empty() {
                cmds.push(cmd);
                cmd = vec![];
            }
            if !token.is_empty() {
                cmd.push(token);
            }
            if ends_with_split && !cmd.is_empty() {
                cmds.push(cmd);
                cmd = vec![];
            }

        }
    }
    if !cmd.is_empty() {
        cmds.push(cmd);
    }
    Some(cmds)
}

/// write tests for parse_line
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_parse_line() {
        assert_eq!(parse_line("echo hello"), Some(vec![vec!["echo".to_string(), "hello".to_string()]]));
        assert_eq!(parse_line("echo hello; echo world"), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()]]));
        assert_eq!(parse_line("echo hello; echo world;"), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()]]));
        assert_eq!(parse_line("echo hello; echo world; "), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()]]));
        assert_eq!(parse_line("echo hello; echo world; echo"), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()], vec!["echo".to_string()]]));
        assert_eq!(parse_line("echo hello ; echo world; echo;"), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()], vec!["echo".to_string()]]));
        assert_eq!(parse_line("echo hello ;echo world; echo ;"), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()], vec!["echo".to_string()]]));
        assert_eq!(parse_line("echo hello; echo world; echo ; "), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()], vec!["echo".to_string()]]));
        assert_eq!(parse_line("echo hello; echo world; echo ; ;"), Some(vec![vec!["echo".to_string(), "hello".to_string()], vec!["echo".to_string(), "world".to_string()], vec!["echo".to_string()]]));
        assert_eq!(parse_line("echo 'hello; yeet'; echo world; echo ; ;"), Some(vec![vec!["echo".to_string(), "hello; yeet".to_string()], vec!["echo".to_string(), "world".to_string()], vec!["echo".to_string()]]));
    }
}