ssh_utils_lib/ssh/
common.rs

1use anyhow::Result;
2use crossterm::terminal::size;
3use russh::{client::Msg, *};
4use std::{convert::TryFrom, time::Duration};
5use std::env;
6use tokio::io::{AsyncReadExt, AsyncWriteExt};
7
8pub struct SshChannel {
9    channel: Channel<Msg>,
10    last_size: (u16, u16),
11}
12
13pub fn default_ssh_config() -> client::Config {
14    client::Config {
15        keepalive_interval: Some(Duration::from_secs(15)),
16        // 其他字段使用默认值
17        ..Default::default()
18    }
19}
20
21impl SshChannel {
22    pub async fn new(channel: Channel<Msg>) -> Result<Self> {
23        let (w, h) = size()?;
24        Ok(Self {
25            channel,
26            last_size: (w, h),
27        })
28    }
29
30    pub async fn call(&mut self, command: &str) -> Result<u32> {
31        let (w, h) = self.last_size;
32
33        // Request an interactive PTY from the server
34        self.channel
35            .request_pty(
36                false,
37                &env::var("TERM").unwrap_or("xterm".into()),
38                w as u32,
39                h as u32,
40                0,
41                0,
42                &[],
43            )
44            .await?;
45        self.channel.exec(true, command).await?;
46
47        let code;
48        let mut stdin = tokio_fd::AsyncFd::try_from(0)?;
49        let mut stdout = tokio_fd::AsyncFd::try_from(1)?;
50        let mut buf = vec![0; 1024];
51        let mut stdin_closed = false;
52
53        loop {
54            tokio::select! {
55                r = stdin.read(&mut buf), if !stdin_closed => {
56                    match r {
57                        Ok(0) => {
58                            stdin_closed = true;
59                            self.channel.eof().await?;
60                        },
61                        Ok(n) => self.channel.data(&buf[..n]).await?,
62                        Err(e) => return Err(e.into()),
63                    };
64                },
65                Some(msg) = self.channel.wait() => {
66                    match msg {
67                        ChannelMsg::Data { ref data } => {
68                            let (w, h) = size()?;
69                            if (w, h) != self.last_size {
70                                self.channel.window_change(w as u32, h as u32, 0, 0).await?;
71                                self.last_size = (w, h);
72                            }
73                            stdout.write_all(data).await?;
74                            stdout.flush().await?;
75                        }
76                        ChannelMsg::ExitStatus { exit_status } => {
77                            code = exit_status;
78                            if !stdin_closed {
79                                self.channel.eof().await?;
80                            }
81                            break;
82                        }
83                        _ => {}
84                    }
85                }
86            }
87        }
88        Ok(code)
89    }
90}