quick_file_transfer/ssh/
remote_session.rs

1use std::{net::ToSocketAddrs, path::Path, time::Duration};
2
3use ssh::{ExecBroker, SessionBroker, SshError, SshResult};
4
5use super::remote_find_free_port::remote_find_free_port;
6
7pub struct RemoteSshSession {
8    session: SessionBroker,
9    latest_executed_cmd: Option<ExecutedCmd>,
10}
11
12impl RemoteSshSession {
13    pub fn new<A>(
14        username: &str,
15        addr: A,
16        timeout: Option<Duration>,
17        private_key_path: Option<&Path>,
18        private_key_dir: Option<&Path>,
19    ) -> Result<Self, SshError>
20    where
21        A: ToSocketAddrs,
22    {
23        let ssh_private_key =
24            crate::ssh::private_key::get_ssh_private_key_path(private_key_path, private_key_dir)
25                .expect("Failed locating SSH private key path");
26        log::trace!("Private key path: {ssh_private_key:?}");
27        let passwd = super::util::get_remote_password_from_env();
28
29        let session = ssh::create_session()
30            .username(username)
31            .password(passwd.as_deref().unwrap_or("root"))
32            .private_key_path(ssh_private_key)
33            .connect_with_timeout(addr, timeout)?;
34        Ok(Self {
35            session: session.run_backend(),
36            latest_executed_cmd: None,
37        })
38    }
39
40    pub fn find_free_port(&mut self, start_port: u16, end_port: u16) -> anyhow::Result<u16> {
41        remote_find_free_port(&mut self.session, start_port, end_port)
42    }
43
44    fn open_run_exec(&mut self, cmd: &str) -> SshResult<ExecutedCmd> {
45        let executed = ExecutedCmd::new(&mut self.session, cmd)?;
46        Ok(executed)
47    }
48
49    /// Runs the remote command through SSH and returns the exit status of the command
50    ///
51    /// If you also want the output of the command, use `run_cmd_get_result`.
52    pub fn run_cmd(&mut self, cmd: &str) -> SshResult<u32> {
53        let executed = self.open_run_exec(cmd)?;
54        let exit_status = executed.exit_status()?;
55        log::trace!("Remote command exit status: {exit_status}");
56        if let Some(terminate_msg) = executed.terminate_msg() {
57            log::trace!("Remote command terminate message: {terminate_msg}");
58        }
59
60        self.latest_executed_cmd = Some(executed);
61
62        Ok(exit_status)
63    }
64
65    /// Runs the remote command through SSH and returns the output/result of the command as UTF-8
66    ///
67    /// This will block until the server closes the channel (meaning the command has to run to end)
68    pub fn run_cmd_get_result(&mut self, cmd: &str) -> SshResult<Vec<u8>> {
69        let mut executed = self.open_run_exec(cmd)?;
70        let exit_status = executed.exit_status()?;
71        log::trace!("Remote command exit status: {exit_status}");
72        if let Some(terminate_msg) = executed.terminate_msg() {
73            log::trace!("Remote command terminate message: {terminate_msg}");
74        }
75
76        let res = executed.results()?;
77        self.latest_executed_cmd = Some(executed);
78        Ok(res)
79    }
80
81    /// If a command has been executed, consumes it and returns its result.
82    ///
83    /// Blocks until the server has closed the connection
84    pub fn get_cmd_output(&mut self) -> Option<Vec<u8>> {
85        let mut exec = self.latest_executed_cmd.take()?;
86        exec.results().ok()
87    }
88
89    /// Close the remote session
90    pub fn close(self) {
91        self.session.close()
92    }
93}
94
95pub struct ExecutedCmd {
96    exec_broker: ExecBroker,
97}
98
99impl ExecutedCmd {
100    pub fn new(session: &mut SessionBroker, cmd: &str) -> SshResult<Self> {
101        let mut exec = session.open_exec()?;
102        exec.send_command(cmd)?;
103        Ok(Self { exec_broker: exec })
104    }
105
106    pub fn exit_status(&self) -> SshResult<u32> {
107        self.exec_broker.exit_status()
108    }
109
110    pub fn terminate_msg(&self) -> Option<String> {
111        let tm = self.exec_broker.terminate_msg().ok()?;
112        if tm.is_empty() {
113            None
114        } else {
115            Some(tm)
116        }
117    }
118
119    /// Blocking
120    pub fn results(&mut self) -> SshResult<Vec<u8>> {
121        self.exec_broker.get_result()
122    }
123}
quick_file_transfer/ssh/remote_session.rs

quick_file_transfer/ssh/
remote_session.rs
