utils_box/connections/
ssh_client.rs

1//! # SSH Client utility
2//! A small SSH utility to connect and manipulate SSH connections to a server.
3//! Useful for executing commands remotely and uploading/downloading files via SSH.
4
5use anyhow::{Result, bail};
6use ssh2::Session;
7use std::{
8    fs::File, io::Read, io::Write, net::TcpStream, os::unix::prelude::PermissionsExt, path::PathBuf,
9};
10
11use crate::{log_info, log_trace};
12
13#[derive(Clone)]
14pub struct SshClient {
15    server_ip: String,
16    server_port: u16,
17    ssh_session: Session,
18}
19
20impl SshClient {
21    pub fn new(
22        server_ip: String,
23        server_port: u16,
24        username: String,
25        password: String,
26    ) -> Result<Self> {
27        // Connect to the local SSH server
28        let tcp_stream = TcpStream::connect(format!("{server_ip}:{server_port}"))?;
29
30        let mut ssh_session = Session::new()?;
31        ssh_session.set_tcp_stream(tcp_stream);
32        ssh_session.handshake()?;
33
34        ssh_session.userauth_password(&username, &password)?;
35
36        if !ssh_session.authenticated() {
37            bail!("[SshClient][new] Authentication FAILED!");
38        }
39
40        Ok(SshClient {
41            server_ip,
42            server_port,
43            ssh_session,
44        })
45    }
46
47    pub fn local(username: String, password: String) -> Result<Self> {
48        Self::new("127.0.0.1".to_string(), 22, username, password)
49    }
50
51    pub fn connection_info(&self) -> (String, u16) {
52        (self.server_ip.clone(), self.server_port)
53    }
54
55    pub fn upload(&self, file: PathBuf, remote_file: PathBuf) -> Result<()> {
56        // Get access to the local file
57        let mut local_file = File::open(&file)?;
58        // Get local permissions, remove group permissions
59        let local_permissions = local_file.metadata()?.permissions().mode() as i32 & 0o777;
60
61        let mut file_stream = vec![];
62        local_file.read_to_end(&mut file_stream)?;
63
64        log_info!(
65            "[SshClient][upload] [{} => {}] Permissions:[{:o}] Size: [{} Bytes]",
66            file.display(),
67            remote_file.display(),
68            local_permissions,
69            file_stream.len(),
70        );
71
72        // Write the file
73        let mut remote_file = self.ssh_session.scp_send(
74            &remote_file,
75            local_permissions,
76            file_stream.len() as u64,
77            None,
78        )?;
79
80        remote_file.write_all(&file_stream)?;
81
82        // Wait for all write operations to finish on remote machine
83        std::thread::sleep(std::time::Duration::from_secs(3));
84
85        remote_file.flush()?;
86
87        remote_file.send_eof()?;
88
89        // Close the channel
90        remote_file.close()?;
91
92        Ok(())
93    }
94
95    pub fn download(&self, remote_file: PathBuf, file: PathBuf) -> Result<()> {
96        // Get access to the local file
97        let mut local = File::create(&file)?;
98
99        // Get access to the remote file
100        let (mut remote, remote_stats) = self.ssh_session.scp_recv(&remote_file)?;
101
102        log_info!(
103            "[SshClient][download] [{} => {}] Permissions:[{:o}] Size: [{} Bytes]",
104            remote_file.display(),
105            file.display(),
106            remote_stats.mode(),
107            remote_stats.size(),
108        );
109
110        // Download content
111        let mut file_stream = vec![];
112        remote.read_to_end(&mut file_stream)?;
113
114        // Close the channel and wait for the whole content to be tranferred
115        remote.send_eof()?;
116        remote.wait_eof()?;
117        remote.close()?;
118        remote.wait_close()?;
119
120        // Write to local
121        local.write_all(&file_stream)?;
122
123        // Wait for all write operations to finish on host
124        std::thread::sleep(std::time::Duration::from_secs(3));
125        local.flush()?;
126
127        // Set permissions
128        local
129            .metadata()?
130            .permissions()
131            .set_mode(remote_stats.mode() as u32);
132
133        Ok(())
134    }
135
136    pub fn execute_cmd(&self, cmd: &str) -> Result<String> {
137        // Open ssh channel
138        let mut channel = self.ssh_session.channel_session()?;
139        channel.exec(cmd)?;
140
141        // Parse STDOUT
142        let mut std_out = String::new();
143        channel.read_to_string(&mut std_out)?;
144        log_trace!("[SshClient][execute_cmd] STDOUT: \n{}", std_out);
145
146        // Close the channel and wait until is confirmed
147        channel.close()?;
148        channel.wait_close()?;
149
150        Ok(std_out)
151    }
152}
utils_box/connections/ssh_client.rs

utils_box/connections/
ssh_client.rs
