Rust 实现ssh操作

使用rust操作ssh,执行命令,然后输出。长时间执行的命令需要一行一行输出。核心是通过BufReader::new来包装一下channel,这种就有了read_lines方法了。

use std::fmt::format;
use std::io::{BufRead, Read};
use std::net::TcpStream;
use std::io::BufReader;

use ssh2::Session;

pub struct Config {
    pub ip: String,
    pub port: i32,
    pub username: String,
    pub password: String,
}

pub struct SSH {
    pub session: Session,
    pub config: Config,
}

impl SSH {
    pub fn new(ip: &str, username: &str, password: &str, port: i32) -> Self {
        if port < 0 || port > 65535 {
            panic!("port{}不合法", port)
        }
        let mut ip = ip.to_string();
        let username = username.to_string();
        let password = password.to_string();
        if !ip.contains(".") {
            ip = format!("##masaike.{}", ip)
        }
        let config = Config { ip, port, username, password };
        let session = Self::connect(&config).unwrap();
        Self { session, config }
    }
    fn connect(c: &Config) -> Result<Session, anyhow::Error> {
        println!("建立对{}:{}的ssh连接", c.ip, c.port);
        let tcp = TcpStream::connect(format!("{}:{}", c.ip, c.port)).unwrap();
        let mut sess = Session::new().unwrap();
        sess.set_tcp_stream(tcp);
        sess.handshake()?;
        sess.userauth_password(c.username.as_str(), c.password.as_str())?;
        println!("连接成功");
        Ok(sess)
    }
    pub fn exec_wait(&self, cmd: &str) {
        println!("执行命令: {}", cmd);
        let mut channel = self.session.channel_session().unwrap();
        channel.exec(cmd).unwrap();
        let mut s = String::new();
        channel.read_to_string(&mut s).unwrap();
        println!("{}", s);
        channel.wait_close();
        println!("响应码:{}", channel.exit_status().unwrap());
    }
    pub fn exec_long(&self, cmd: &str) {
        println!("执行长时间命令: {}", cmd);
        let mut channel = self.session.channel_session().unwrap();
        channel.exec(cmd).unwrap();
        let mut reader = BufReader::new(&mut channel);
        loop{
            let mut line = String::new();
            if let Ok(s) = reader.read_line(&mut line){
                if s>0{
                    print!("{}",line);
                }else{
                    break
                }
            }else{
                break
            }
        }
        channel.wait_close();
        println!("响应码:{}", channel.exit_status().unwrap());
    }
}

#[cfg(test)]
mod tests {
    // 注意这个惯用法:在 tests 模块中,从外部作用域导入所有名字。
    use super::*;

    #[test]
    #[should_panic]
    fn test_cannot_connect() {
        let _ = SSH::new("masaike", "root", "", 22);
    }

    #[test]
    fn test_exec_wait() {
        let ssh = SSH::new("masaike", "root", "balabala", 22);
        ssh.exec_wait("ls")
    }

    #[test]
    fn test_exec_long() {
        let ssh = SSH::new("masaike", "root", "masaike", 22);
        ssh.exec_long("echo ok && cd /code/gitlab/masaike && sh make.sh web")
    }
}