devops_armory/rustible/vm/client/
client.rs

1use std::path::Path;
2use std::sync::Arc;
3use std::time::Duration;
4
5use anyhow::Result;
6use async_trait::async_trait;
7use russh::keys::*;
8use russh::*;
9use tokio::net::ToSocketAddrs;
10use std::io::stdout;
11use std::io::Write;
12
13pub struct Client {}
14
15// More SSH event handlers
16// can be defined in this trait
17// In this example, we're only using Channel, so these aren't needed.
18#[async_trait]
19impl client::Handler for Client {
20    type Error = russh::Error;
21
22    async fn check_server_key(
23        &mut self,
24        _server_public_key: &key::PublicKey,
25    ) -> Result<bool, Self::Error> {
26        Ok(true)
27    }
28}
29
30pub struct Session {
31    session: client::Handle<Client>,
32}
33
34impl Session {
35    pub async fn connect<P: AsRef<Path>, A: ToSocketAddrs>(
36        key_path: P,
37        user: impl Into<String>,
38        addrs: A,
39    ) -> Result<Self> {
40        let key_pair = load_secret_key(key_path, None)?;
41        let config = client::Config {
42            inactivity_timeout: Some(Duration::from_secs(5)),
43            ..<_>::default()
44        };
45
46        let config = Arc::new(config);
47        let sh = Client {};
48
49        let mut session = client::connect(config, addrs, sh).await?;
50        let auth_res = session
51            .authenticate_publickey(user, Arc::new(key_pair))
52            .await?;
53
54        if !auth_res {
55            anyhow::bail!("Authentication failed");
56        }
57
58        Ok(Self { session })
59    }
60
61    pub async fn call(&mut self, command: &str) -> Result<u32> {
62        let mut channel = self.session.channel_open_session().await?;
63        channel.exec(true, command).await?;
64
65        let mut code = None;
66        let mut stdout = stdout();
67
68        loop {
69            // There's an event available on the session channel
70            let Some(msg) = channel.wait().await else {
71                break;
72            };
73            match msg {
74                // Write data to the terminal
75                ChannelMsg::Data { ref data } => {
76                    stdout.write_all(data)?;
77                    stdout.flush()?;
78                }
79                // The command has returned an exit code
80                ChannelMsg::ExitStatus { exit_status } => {
81                    code = Some(exit_status);
82                    // cannot leave the loop immediately, there might still be more data to receive
83                }
84                _ => {}
85            }
86        }
87        Ok(code.expect("program did not exit cleanly"))
88    }
89
90    pub async fn close(&mut self) -> Result<()> {
91        self.session
92            .disconnect(Disconnect::ByApplication, "", "English")
93            .await?;
94        Ok(())
95    }
96}
devops_armory/rustible/vm/client/client.rs

devops_armory/rustible/vm/client/
client.rs
