lmrc_ssh/

client.rs

//! SSH client implementation.

use crate::{CommandOutput, Error, Result};
use ssh2::Session;
use std::io::Read;
use std::net::TcpStream;
use std::path::Path;

/// Authentication methods for SSH connections.
#[derive(Debug, Clone)]
pub enum AuthMethod {
    /// Authenticate using a username and password.
    Password {
        /// The username for authentication.
        username: String,
        /// The password for authentication.
        password: String,
    },

    /// Authenticate using a public/private key pair.
    PublicKey {
        /// The username for authentication.
        username: String,
        /// Path to the private key file.
        private_key_path: String,
        /// Optional passphrase for the private key.
        passphrase: Option<String>,
    },
}

/// An SSH client for executing remote commands.
///
/// # Examples
///
/// ```rust,no_run
/// use lmrc_ssh::{SshClient, AuthMethod};
///
/// # fn main() -> Result<(), lmrc_ssh::Error> {
/// let mut client = SshClient::new("example.com", 22)?
///     .with_auth(AuthMethod::Password {
///         username: "user".to_string(),
///         password: "pass".to_string(),
///     })
///     .connect()?;
///
/// let output = client.execute("hostname")?;
/// println!("Hostname: {}", output.stdout);
/// # Ok(())
/// # }
/// ```
pub struct SshClient {
    host: String,
    port: u16,
    auth: Option<AuthMethod>,
    session: Option<Session>,
}

impl std::fmt::Debug for SshClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SshClient")
            .field("host", &self.host)
            .field("port", &self.port)
            .field("auth", &self.auth)
            .field("connected", &self.is_connected())
            .finish()
    }
}

impl SshClient {
    /// Creates a new SSH client instance.
    ///
    /// # Arguments
    ///
    /// * `host` - The hostname or IP address of the remote server
    /// * `port` - The SSH port (typically 22)
    ///
    /// # Examples
    ///
    /// ```rust
    /// use lmrc_ssh::SshClient;
    ///
    /// # fn main() -> Result<(), lmrc_ssh::Error> {
    /// let client = SshClient::new("example.com", 22)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn new(host: impl Into<String>, port: u16) -> Result<Self> {
        let host = host.into();

        if host.is_empty() {
            return Err(Error::InvalidConfig("Host cannot be empty".to_string()));
        }

        if port == 0 {
            return Err(Error::InvalidConfig("Port cannot be 0".to_string()));
        }

        Ok(Self {
            host,
            port,
            auth: None,
            session: None,
        })
    }

    /// Sets the authentication method for the connection.
    ///
    /// # Arguments
    ///
    /// * `auth` - The authentication method to use
    ///
    /// # Examples
    ///
    /// ```rust
    /// use lmrc_ssh::{SshClient, AuthMethod};
    ///
    /// # fn main() -> Result<(), lmrc_ssh::Error> {
    /// let client = SshClient::new("example.com", 22)?
    ///     .with_auth(AuthMethod::Password {
    ///         username: "user".to_string(),
    ///         password: "pass".to_string(),
    ///     });
    /// # Ok(())
    /// # }
    /// ```
    pub fn with_auth(mut self, auth: AuthMethod) -> Self {
        self.auth = Some(auth);
        self
    }

    /// Establishes the SSH connection and authenticates.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The TCP connection fails
    /// - The SSH handshake fails
    /// - Authentication fails
    /// - No authentication method was set
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use lmrc_ssh::{SshClient, AuthMethod};
    ///
    /// # fn main() -> Result<(), lmrc_ssh::Error> {
    /// let mut client = SshClient::new("example.com", 22)?
    ///     .with_auth(AuthMethod::Password {
    ///         username: "user".to_string(),
    ///         password: "pass".to_string(),
    ///     })
    ///     .connect()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn connect(mut self) -> Result<Self> {
        let auth = self
            .auth
            .as_ref()
            .ok_or_else(|| Error::InvalidConfig("No authentication method set".to_string()))?;

        let tcp = TcpStream::connect(format!("{}:{}", self.host, self.port)).map_err(|e| {
            Error::ConnectionFailed {
                host: self.host.clone(),
                port: self.port,
                source: e,
            }
        })?;

        let mut session = Session::new()?;
        session.set_tcp_stream(tcp);
        session.handshake()?;

        match auth {
            AuthMethod::Password { username, password } => {
                session.userauth_password(username, password).map_err(|e| {
                    Error::AuthenticationFailed {
                        username: username.clone(),
                        reason: e.to_string(),
                    }
                })?;
            }
            AuthMethod::PublicKey {
                username,
                private_key_path,
                passphrase,
            } => {
                let key_path = Path::new(private_key_path);

                if !key_path.exists() {
                    return Err(Error::PrivateKeyNotFound {
                        path: private_key_path.clone(),
                    });
                }

                session
                    .userauth_pubkey_file(username, None, key_path, passphrase.as_deref())
                    .map_err(|e| Error::AuthenticationFailed {
                        username: username.clone(),
                        reason: e.to_string(),
                    })?;
            }
        }

        if !session.authenticated() {
            let username = match auth {
                AuthMethod::Password { username, .. } => username.clone(),
                AuthMethod::PublicKey { username, .. } => username.clone(),
            };
            return Err(Error::AuthenticationFailed {
                username,
                reason: "Authentication completed but session is not authenticated".to_string(),
            });
        }

        self.session = Some(session);
        Ok(self)
    }

    /// Executes a command on the remote server.
    ///
    /// # Arguments
    ///
    /// * `command` - The command to execute
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The client is not connected
    /// - Failed to open an SSH channel
    /// - Failed to execute the command
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use lmrc_ssh::{SshClient, AuthMethod};
    ///
    /// # fn main() -> Result<(), lmrc_ssh::Error> {
    /// let mut client = SshClient::new("example.com", 22)?
    ///     .with_auth(AuthMethod::Password {
    ///         username: "user".to_string(),
    ///         password: "pass".to_string(),
    ///     })
    ///     .connect()?;
    ///
    /// let output = client.execute("ls -la /tmp")?;
    /// println!("Files:\n{}", output.stdout);
    /// # Ok(())
    /// # }
    /// ```
    pub fn execute(&mut self, command: &str) -> Result<CommandOutput> {
        let session = self.session.as_ref().ok_or(Error::NotConnected)?;

        let mut channel = session
            .channel_session()
            .map_err(|e| Error::ChannelFailed(e.to_string()))?;

        channel
            .exec(command)
            .map_err(|e| Error::ExecutionFailed(e.to_string()))?;

        let mut stdout = String::new();
        channel
            .read_to_string(&mut stdout)
            .map_err(|e| Error::ExecutionFailed(format!("Failed to read stdout: {}", e)))?;

        let mut stderr = String::new();
        channel
            .stderr()
            .read_to_string(&mut stderr)
            .map_err(|e| Error::ExecutionFailed(format!("Failed to read stderr: {}", e)))?;

        channel
            .wait_close()
            .map_err(|e| Error::ExecutionFailed(format!("Failed to close channel: {}", e)))?;

        let exit_status = channel.exit_status()?;

        Ok(CommandOutput::new(stdout, stderr, exit_status))
    }

    /// Executes multiple commands sequentially.
    ///
    /// # Arguments
    ///
    /// * `commands` - A slice of commands to execute
    ///
    /// # Returns
    ///
    /// Returns a vector of `CommandOutput` for each executed command.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use lmrc_ssh::{SshClient, AuthMethod};
    ///
    /// # fn main() -> Result<(), lmrc_ssh::Error> {
    /// let mut client = SshClient::new("example.com", 22)?
    ///     .with_auth(AuthMethod::Password {
    ///         username: "user".to_string(),
    ///         password: "pass".to_string(),
    ///     })
    ///     .connect()?;
    ///
    /// let outputs = client.execute_batch(&["whoami", "hostname", "pwd"])?;
    /// for output in outputs {
    ///     println!("{}", output.stdout);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn execute_batch(&mut self, commands: &[&str]) -> Result<Vec<CommandOutput>> {
        commands.iter().map(|cmd| self.execute(cmd)).collect()
    }

    /// Returns the hostname this client is configured to connect to.
    pub fn host(&self) -> &str {
        &self.host
    }

    /// Returns the port this client is configured to connect to.
    pub fn port(&self) -> u16 {
        self.port
    }

    /// Returns whether the client is currently connected.
    pub fn is_connected(&self) -> bool {
        self.session.is_some()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_new_client() {
        let client = SshClient::new("example.com", 22).unwrap();
        assert_eq!(client.host(), "example.com");
        assert_eq!(client.port(), 22);
        assert!(!client.is_connected());
    }

    #[test]
    fn test_empty_host() {
        let result = SshClient::new("", 22);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidConfig(_)));
    }

    #[test]
    fn test_zero_port() {
        let result = SshClient::new("example.com", 0);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidConfig(_)));
    }

    #[test]
    fn test_with_auth() {
        let client = SshClient::new("example.com", 22)
            .unwrap()
            .with_auth(AuthMethod::Password {
                username: "user".to_string(),
                password: "pass".to_string(),
            });
        assert!(client.auth.is_some());
    }

    #[test]
    fn test_connect_without_auth() {
        let client = SshClient::new("example.com", 22).unwrap();
        let result = client.connect();
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidConfig(_)));
    }
}