ssh2-config-rs 0.7.2

//! # client
//!
//! Ssh2-config implementation with a russh client

use std::env::args;
use std::fs::File;
use std::io::BufReader;
use std::net::{SocketAddr, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;
use std::time::Duration;

use dirs::home_dir;
use russh::client::{AuthResult, Config, Handler};
use russh::keys::{Algorithm, PublicKey};
use russh::*;
use ssh2_config_rs::{HostParams, ParseRule, SshConfig};
use std::borrow::Cow;
use tokio::net::TcpStream;
use tokio::time::timeout;

struct ClientHandler;

impl Handler for ClientHandler {
    type Error = russh::Error;

    async fn check_server_key(
        &mut self,
        _server_public_key: &PublicKey,
    ) -> Result<bool, Self::Error> {
        // Accept all server keys for this example
        // In production, you should verify the server key
        Ok(true)
    }

    async fn auth_banner(
        &mut self,
        banner: &str,
        _session: &mut russh::client::Session,
    ) -> Result<(), Self::Error> {
        println!("{}", banner);
        Ok(())
    }
}

#[tokio::main]
async fn main() {
    // get args
    let args: Vec<String> = args().collect();
    let address = match args.get(1) {
        Some(addr) => addr.to_string(),
        None => {
            usage();
            exit(255)
        }
    };
    // check path
    let config_path = match args.get(2) {
        Some(p) => PathBuf::from(p),
        None => {
            let mut p = home_dir().expect("Failed to get home_dir for guest OS");
            p.extend(Path::new(".ssh/config"));
            p
        }
    };
    // Open config file
    let config = read_config(config_path.as_path());
    let params = config.query(address.as_str());

    if let Err(e) = connect(address.as_str(), &params).await {
        eprintln!("Connection failed: {}", e);
        exit(1);
    }
}

fn usage() {
    eprintln!("Usage: cargo run --example client -- <address:port> [config-path]");
}

fn read_config(p: &Path) -> SshConfig {
    let mut reader = match File::open(p) {
        Ok(f) => BufReader::new(f),
        Err(err) => panic!("Could not open file '{}': {}", p.display(), err),
    };
    match SshConfig::default().parse(&mut reader, ParseRule::STRICT) {
        Ok(config) => config,
        Err(err) => panic!("Failed to parse configuration: {}", err),
    }
}

async fn connect(host: &str, params: &HostParams) -> Result<(), russh::Error> {
    // Resolve host
    let host = match params.host_name.as_deref() {
        Some(h) => h,
        None => host,
    };
    let port = params.port.unwrap_or(22);
    let host = match host.contains(':') {
        true => host.to_string(),
        false => format!("{}:{}", host, port),
    };
    println!("Connecting to host {}...", host);

    let socket_addresses: Vec<SocketAddr> = match host.to_socket_addrs() {
        Ok(s) => s.collect(),
        Err(err) => {
            panic!("Could not parse host: {}", err);
        }
    };

    // Try addresses
    let mut stream: Option<TcpStream> = None;
    let connect_timeout = params.connect_timeout.unwrap_or(Duration::from_secs(30));

    for socket_addr in socket_addresses.iter() {
        match timeout(connect_timeout, TcpStream::connect(socket_addr)).await {
            Ok(Ok(s)) => {
                println!("Established connection with {}", socket_addr);
                stream = Some(s);
                break;
            }
            Ok(Err(_)) => continue,
            Err(_) => continue, // timeout
        }
    }

    let stream = stream.ok_or_else(|| russh::Error::ConnectionTimeout)?;

    // Configure session
    let mut config = Config::default();
    configure_session(&mut config, params);

    // Connect using the manually created stream
    let config = Arc::new(config);
    let handler = ClientHandler;

    let mut handle = russh::client::connect_stream(config, stream, handler).await?;

    // Get username
    let username = match params.user.as_ref() {
        Some(u) => {
            println!("Using username '{}'", u);
            u.clone()
        }
        None => read_secret("Username: "),
    };

    // Authenticate with password
    let password = read_secret("Password: ");
    match handle.authenticate_password(username, password).await? {
        AuthResult::Success => {
            println!("Authentication successful!");
        }
        AuthResult::Failure { .. } => {
            return Err(russh::Error::NotAuthenticated);
        }
    }

    println!("Connection OK!");

    // Disconnect
    handle
        .disconnect(Disconnect::ByApplication, "mandi mandi!", "en")
        .await?;

    Ok(())
}

fn configure_session(config: &mut Config, params: &HostParams) {
    println!("Configuring session...");

    // Compression
    if let Some(compress) = params.compression {
        println!("compression: {}", compress);
        // Note: russh compression is configured via Preferred, which is set below
        // Compression algorithm names would need to be mapped if specified
    }

    // Keepalive
    if params.tcp_keep_alive.unwrap_or(false) && params.server_alive_interval.is_some() {
        let interval = params.server_alive_interval.unwrap();
        println!("keepalive interval: {} seconds", interval.as_secs());
        config.keepalive_interval = Some(interval);
    }

    // Configure preferred algorithms
    // Convert algorithm strings to russh Name types
    let mut preferred = Preferred::default();

    // KEX algorithms
    let kex_algorithms: Vec<kex::Name> = params
        .kex_algorithms
        .algorithms()
        .iter()
        .filter_map(|s| kex::Name::try_from(s.as_str()).ok())
        .collect();
    if !kex_algorithms.is_empty() {
        println!(
            "KEX algorithms: {:?}",
            kex_algorithms
                .iter()
                .map(|n| n.as_ref())
                .collect::<Vec<_>>()
        );
        preferred.kex = Cow::Owned(kex_algorithms);
    }

    // Host key algorithms
    let host_key_algorithms: Vec<Algorithm> = params
        .host_key_algorithms
        .algorithms()
        .iter()
        .filter_map(|s| Algorithm::new(s.as_str()).ok())
        .collect();
    if !host_key_algorithms.is_empty() {
        println!(
            "Host key algorithms: {:?}",
            host_key_algorithms
                .iter()
                .map(|a| a.as_ref())
                .collect::<Vec<_>>()
        );
        preferred.key = Cow::Owned(host_key_algorithms);
    }

    // Ciphers
    let ciphers: Vec<cipher::Name> = params
        .ciphers
        .algorithms()
        .iter()
        .filter_map(|s| cipher::Name::try_from(s.as_str()).ok())
        .collect();
    if !ciphers.is_empty() {
        println!(
            "Ciphers: {:?}",
            ciphers.iter().map(|n| n.as_ref()).collect::<Vec<_>>()
        );
        preferred.cipher = Cow::Owned(ciphers);
    }

    // MACs
    let macs: Vec<mac::Name> = params
        .mac
        .algorithms()
        .iter()
        .filter_map(|s| mac::Name::try_from(s.as_str()).ok())
        .collect();
    if !macs.is_empty() {
        println!(
            "MACs: {:?}",
            macs.iter().map(|n| n.as_ref()).collect::<Vec<_>>()
        );
        preferred.mac = Cow::Owned(macs);
    }

    config.preferred = preferred;
}

fn read_secret(prompt: &str) -> String {
    rpassword::prompt_password(prompt).expect("Failed to read from stdin")
}