use std::net::IpAddr; use std::time::Duration; use std::time::Instant; use crate::BrocadeOptions; use crate::Error; use crate::ExecutionMode; use crate::TimeoutConfig; use crate::ssh; use log::debug; use log::info; use russh::ChannelMsg; use tokio::time::timeout; #[derive(Debug)] pub struct BrocadeShell { ip: IpAddr, port: u16, username: String, password: String, options: BrocadeOptions, before_all_commands: Vec, after_all_commands: Vec, } impl BrocadeShell { pub async fn init( ip_addresses: &[IpAddr], port: u16, username: &str, password: &str, options: Option, ) -> Result { let ip = ip_addresses .first() .ok_or_else(|| Error::ConfigurationError("No IP addresses provided".to_string()))?; let base_options = options.unwrap_or_default(); let options = ssh::try_init_client(username, password, ip, base_options).await?; Ok(Self { ip: *ip, port, username: username.to_string(), password: password.to_string(), before_all_commands: vec![], after_all_commands: vec![], options, }) } pub async fn open_session(&self, mode: ExecutionMode) -> Result { BrocadeSession::open( self.ip, self.port, &self.username, &self.password, self.options.clone(), mode, ) .await } pub async fn with_session(&self, mode: ExecutionMode, callback: F) -> Result where F: FnOnce( &mut BrocadeSession, ) -> std::pin::Pin< Box> + Send + '_>, >, { let mut session = self.open_session(mode).await?; let _ = session.run_commands(self.before_all_commands.clone()).await; let result = callback(&mut session).await; let _ = session.run_commands(self.after_all_commands.clone()).await; session.close().await?; result } pub async fn run_command(&self, command: &str, mode: ExecutionMode) -> Result { let mut session = self.open_session(mode).await?; let _ = session.run_commands(self.before_all_commands.clone()).await; let result = session.run_command(command).await; let _ = session.run_commands(self.after_all_commands.clone()).await; session.close().await?; result } pub async fn run_commands( &self, commands: Vec, mode: ExecutionMode, ) -> Result<(), Error> { let mut session = self.open_session(mode).await?; let _ = session.run_commands(self.before_all_commands.clone()).await; let result = session.run_commands(commands).await; let _ = session.run_commands(self.after_all_commands.clone()).await; session.close().await?; result } pub fn before_all(&mut self, commands: Vec) { self.before_all_commands = commands; } pub fn after_all(&mut self, commands: Vec) { self.after_all_commands = commands; } } pub struct BrocadeSession { pub channel: russh::Channel, pub mode: ExecutionMode, pub options: BrocadeOptions, } impl BrocadeSession { pub async fn open( ip: IpAddr, port: u16, username: &str, password: &str, options: BrocadeOptions, mode: ExecutionMode, ) -> Result { let client = ssh::create_client(ip, port, username, password, &options).await?; let mut channel = client.channel_open_session().await?; channel .request_pty(false, "vt100", 80, 24, 0, 0, &[]) .await?; channel.request_shell(false).await?; wait_for_shell_ready(&mut channel, &options.timeouts).await?; if let ExecutionMode::Privileged = mode { try_elevate_session(&mut channel, username, password, &options.timeouts).await?; } Ok(Self { channel, mode, options, }) } pub async fn close(&mut self) -> Result<(), Error> { debug!("[Brocade] Closing session..."); self.channel.data(&b"exit\n"[..]).await?; if let ExecutionMode::Privileged = self.mode { self.channel.data(&b"exit\n"[..]).await?; } let start = Instant::now(); while start.elapsed() < self.options.timeouts.cleanup { match timeout(self.options.timeouts.message_wait, self.channel.wait()).await { Ok(Some(ChannelMsg::Close)) => break, Ok(Some(_)) => continue, Ok(None) | Err(_) => break, } } debug!("[Brocade] Session closed."); Ok(()) } pub async fn run_command(&mut self, command: &str) -> Result { if self.should_skip_command(command) { return Ok(String::new()); } debug!("[Brocade] Running command: '{command}'..."); self.channel .data(format!("{}\n", command).as_bytes()) .await?; tokio::time::sleep(Duration::from_millis(100)).await; let output = self.collect_command_output().await?; let output = String::from_utf8(output) .map_err(|_| Error::UnexpectedError("Invalid UTF-8 in command output".to_string()))?; self.check_for_command_errors(&output, command)?; Ok(output) } pub async fn run_commands(&mut self, commands: Vec) -> Result<(), Error> { for command in commands { self.run_command(&command).await?; } Ok(()) } fn should_skip_command(&self, command: &str) -> bool { if (command.starts_with("write") || command.starts_with("deploy")) && self.options.dry_run { info!("[Brocade] Dry-run mode enabled, skipping command: {command}"); return true; } false } async fn collect_command_output(&mut self) -> Result, Error> { let mut output = Vec::new(); let start = Instant::now(); let read_timeout = Duration::from_millis(500); let log_interval = Duration::from_secs(3); let mut last_log = Instant::now(); loop { if start.elapsed() > self.options.timeouts.command_execution { return Err(Error::TimeoutError( "Timeout waiting for command completion.".into(), )); } if start.elapsed() > Duration::from_secs(5) && last_log.elapsed() > log_interval { info!("[Brocade] Waiting for command output..."); last_log = Instant::now(); } match timeout(read_timeout, self.channel.wait()).await { Ok(Some(ChannelMsg::Data { data } | ChannelMsg::ExtendedData { data, .. })) => { output.extend_from_slice(&data); let current_output = String::from_utf8_lossy(&output); if current_output.contains('>') || current_output.contains('#') { return Ok(output); } } Ok(Some(ChannelMsg::Eof | ChannelMsg::Close)) => return Ok(output), Ok(Some(ChannelMsg::ExitStatus { exit_status })) => { debug!("[Brocade] Command exit status: {exit_status}"); } Ok(Some(_)) => continue, Ok(None) | Err(_) => { if output.is_empty() { if let Ok(None) = timeout(read_timeout, self.channel.wait()).await { break; } continue; } tokio::time::sleep(Duration::from_millis(100)).await; let current_output = String::from_utf8_lossy(&output); if current_output.contains('>') || current_output.contains('#') { return Ok(output); } } } } Ok(output) } fn check_for_command_errors(&self, output: &str, command: &str) -> Result<(), Error> { const ERROR_PATTERNS: &[&str] = &[ "invalid input", "syntax error", "command not found", "unknown command", "permission denied", "access denied", "authentication failed", "configuration error", "failed to", "error:", ]; let output_lower = output.to_lowercase(); if ERROR_PATTERNS.iter().any(|&p| output_lower.contains(p)) { return Err(Error::CommandError(format!( "Command '{command}' failed: {}", output.trim() ))); } if !command.starts_with("show") && output.trim().is_empty() { return Err(Error::CommandError(format!( "Command '{command}' produced no output" ))); } Ok(()) } } async fn wait_for_shell_ready( channel: &mut russh::Channel, timeouts: &TimeoutConfig, ) -> Result<(), Error> { let mut buffer = Vec::new(); let start = Instant::now(); while start.elapsed() < timeouts.shell_ready { match timeout(timeouts.message_wait, channel.wait()).await { Ok(Some(ChannelMsg::Data { data })) => { buffer.extend_from_slice(&data); let output = String::from_utf8_lossy(&buffer); let output = output.trim(); if output.ends_with('>') || output.ends_with('#') { debug!("[Brocade] Shell ready"); return Ok(()); } } Ok(Some(_)) => continue, Ok(None) => break, Err(_) => continue, } } Ok(()) } async fn try_elevate_session( channel: &mut russh::Channel, username: &str, password: &str, timeouts: &TimeoutConfig, ) -> Result<(), Error> { channel.data(&b"enable\n"[..]).await?; let start = Instant::now(); let mut buffer = Vec::new(); while start.elapsed() < timeouts.shell_ready { match timeout(timeouts.message_wait, channel.wait()).await { Ok(Some(ChannelMsg::Data { data })) => { buffer.extend_from_slice(&data); let output = String::from_utf8_lossy(&buffer); if output.ends_with('#') { debug!("[Brocade] Privileged mode established"); return Ok(()); } if output.contains("User Name:") { channel.data(format!("{}\n", username).as_bytes()).await?; buffer.clear(); } else if output.contains("Password:") { channel.data(format!("{}\n", password).as_bytes()).await?; buffer.clear(); } else if output.contains('>') { return Err(Error::AuthenticationError( "Enable authentication failed".into(), )); } } Ok(Some(_)) => continue, Ok(None) => break, Err(_) => continue, } } let output = String::from_utf8_lossy(&buffer); if output.ends_with('#') { debug!("[Brocade] Privileged mode established"); Ok(()) } else { Err(Error::AuthenticationError(format!( "Enable failed. Output:\n{output}" ))) } }