t_ssh_client/
client.rs

1use std::io::Write;
2use std::net::ToSocketAddrs;
3use std::sync::Arc;
4use std::time::Duration;
5
6use tokio::{fs, time};
7
8use crate::error::ClientError;
9use crate::{Handler, Output};
10
11pub enum AuthMethod {
12    // Password string
13    Password(String),
14    // Secret key path
15    Key(String),
16}
17
18pub struct ClientBuilder {
19    username: String,
20    auth: Option<AuthMethod>,
21    connect_timeout: Duration,
22}
23
24impl ClientBuilder {
25    pub fn new() -> Self {
26        Self {
27            username: String::default(),
28            auth: None,
29            connect_timeout: Duration::from_secs(10),
30        }
31    }
32
33    pub fn username<S: ToString>(&mut self, username: S) -> &mut Self {
34        self.username = username.to_string();
35        self
36    }
37
38    pub fn auth(&mut self, auth: AuthMethod) -> &mut Self {
39        self.auth = Some(auth);
40        self
41    }
42
43    pub fn connect_timeout(&mut self, timeout: Duration) -> &mut Self {
44        self.connect_timeout = timeout;
45        self
46    }
47
48    pub async fn connect<T: ToSocketAddrs>(&self, addr: T) -> Result<Client, ClientError> {
49        let config = Arc::new(thrussh::client::Config::default());
50        match time::timeout(
51            self.connect_timeout,
52            thrussh::client::connect(config, addr, Handler::default()),
53        )
54        .await
55        {
56            Ok(Ok(handle)) => {
57                if self.username.is_empty() {
58                    return Err(ClientError::UsernameEmpty);
59                }
60                let mut client = Client {
61                    inner: handle,
62                    username: self.username.clone(),
63                };
64                match &self.auth {
65                    Some(AuthMethod::Password(pass)) => client.auth_with_password(&pass).await?,
66                    Some(AuthMethod::Key(path)) => {
67                        let secret: String = fs::read_to_string(path).await?;
68                        let key_pair = thrussh_keys::decode_secret_key(&secret, None)?;
69                        client.auth_with_key_pair(Arc::new(key_pair)).await?
70                    }
71                    None => {}
72                }
73                Ok(client)
74            }
75            Ok(Err(err)) => return Err(err),
76            Err(_) => return Err(ClientError::Timeout),
77        }
78    }
79}
80
81pub struct Client {
82    username: String,
83    inner: thrussh::client::Handle<Handler>,
84}
85
86impl Client {
87    pub fn builder() -> ClientBuilder {
88        ClientBuilder::new()
89    }
90
91    pub(crate) async fn auth_with_password(&mut self, password: &str) -> Result<(), ClientError> {
92        match self
93            .inner
94            .authenticate_password(&self.username, password)
95            .await
96        {
97            Ok(true) => Ok(()),
98            Ok(false) => Err(ClientError::AuthFailed(String::from(
99                "username or password is wrong!",
100            ))),
101            Err(e) => Err(ClientError::ClientFailed(e)),
102        }
103    }
104
105    pub(crate) async fn auth_with_key_pair(
106        &mut self,
107        key: Arc<thrussh_keys::key::KeyPair>,
108    ) -> Result<(), ClientError> {
109        match self.inner.authenticate_publickey(&self.username, key).await {
110            Ok(true) => Ok(()),
111            Ok(false) => Err(ClientError::AuthFailed(String::from(
112                "username or key is wrong!",
113            ))),
114            Err(e) => Err(ClientError::ClientFailed(e)),
115        }
116    }
117
118    #[allow(unused_variables)]
119    pub async fn output(&mut self, command: &str) -> Result<Output, ClientError> {
120        let mut channel = self.inner.channel_open_session().await?;
121        channel.exec(true, command).await?;
122        let mut res = Output::default();
123        while let Some(msg) = channel.wait().await {
124            match msg {
125                thrussh::ChannelMsg::Data { ref data } => {
126                    res.stdout.write_all(&data)?;
127                }
128                thrussh::ChannelMsg::ExtendedData { ref data, ext } => {
129                    res.stderr.write_all(&data)?;
130                }
131                thrussh::ChannelMsg::ExitStatus { exit_status } => {
132                    res.code = Some(exit_status);
133                }
134                _ => {}
135            }
136        }
137        Ok(res)
138    }
139
140    pub async fn close(&mut self) -> Result<(), ClientError> {
141        self.inner
142            .disconnect(thrussh::Disconnect::ByApplication, "", "English")
143            .await?;
144        Ok(())
145    }
146}
t_ssh_client/client.rs

t_ssh_client/
client.rs
