harmony/brocade/src/shell.rs

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<String>,
    after_all_commands: Vec<String>,
}

impl BrocadeShell {
    pub async fn init(
        ip_addresses: &[IpAddr],
        port: u16,
        username: &str,
        password: &str,
        options: Option<BrocadeOptions>,
    ) -> Result<Self, Error> {
        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, Error> {
        BrocadeSession::open(
            self.ip,
            self.port,
            &self.username,
            &self.password,
            self.options.clone(),
            mode,
        )
        .await
    }

    pub async fn with_session<F, R>(&self, mode: ExecutionMode, callback: F) -> Result<R, Error>
    where
        F: FnOnce(
            &mut BrocadeSession,
        ) -> std::pin::Pin<
            Box<dyn std::future::Future<Output = Result<R, Error>> + 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<String, Error> {
        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<String>,
        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<String>) {
        self.before_all_commands = commands;
    }

    pub fn after_all(&mut self, commands: Vec<String>) {
        self.after_all_commands = commands;
    }
}

pub struct BrocadeSession {
    pub channel: russh::Channel<russh::client::Msg>,
    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<Self, Error> {
        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<String, Error> {
        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<String>) -> 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<Vec<u8>, Error> {
        let mut output = Vec::new();
        let start = Instant::now();
        let read_timeout = Duration::from_millis(500);
        let log_interval = Duration::from_secs(5);
        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() > self.options.timeouts.command_output
                && 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 error: {}",
                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<russh::client::Msg>,
    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<russh::client::Msg>,
    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}"
        )))
    }
}