multi-ssh/src/shell_interface.rs

use crate::logger::{LogLevel, Logger};
use crate::server::{Server, ServerAddress};
use crate::{log, osf};
use std::error::Error;
use std::ffi::OsString;
use std::fmt::{Debug, Display, Formatter};
use std::iter::once;
use std::path::{Path, PathBuf};
use std::process::{Command, Output};
use std::{io, process};

#[derive(Debug)]
pub struct EnvCommand<'a, E> {
  command: ShellCommand,
  environment: &'a mut E,
}

impl<E> EnvCommand<'_, E>
where
  E: ShellInterface,
{
  pub fn run(self) -> CommandResult<ExitStatus, StartError> {
    self.environment.run_command(self.command)
  }
  
  pub fn output(self) -> CommandResult<CommandOutput, StartError> {
    self.environment.collect_command_output(self.command)
  }
  
  pub fn run_logged(self, logger: &Logger) -> CommandResult<LoggedRunOutput, StartError>
  where
    Self: Sized,
  {
    match logger.level {
      LogLevel::Debug | LogLevel::Info => {
        let res = self.run();
        CommandResult {
          result: res.result.map(LoggedRunOutput::from),
          command: res.command,
        }
      }
      LogLevel::Error => {
        let res = self.output();
        CommandResult {
          result: res.result.map(LoggedRunOutput::from),
          command: res.command,
        }
      }
    }
  }
}

#[derive(Debug, Clone)]
pub enum ShellCommand {
  Ssh {
    address: String,
    server_command: ServerCommand,
  },
  Scp {
    source: ScpParam,
    destination: ScpParam,
  },
  SshAgent,
  ShhAdd,
  Editor(Vec<OsString>),
  Execute {
    working_directory: PathBuf,
    command: Vec<OsString>,
  },
}

impl ShellCommand {
  pub fn in_env<E>(self, environment: &mut E) -> EnvCommand<E> {
    EnvCommand {
      command: self,
      environment,
    }
  }
}

#[derive(Debug, Clone)]
pub enum ServerCommand {
  Realpath {
    path: PathBuf,
  },
  Ls {
    dir: PathBuf,
  },
  Rm {
    file: PathBuf,
  },
  Mv {
    source: PathBuf,
    destination: PathBuf,
  },
  Execute {
    working_directory: PathBuf,
    command: OsString,
  },
}

#[derive(Debug, Clone)]
pub struct ScpParam {
  pub server: Option<String>,
  pub path: PathBuf,
}

impl<P> From<(&Server, P)> for ScpParam
where
  P: AsRef<Path>,
{
  fn from((server, path): (&Server, P)) -> Self {
    Self {
      server: match &server.address {
        ServerAddress::Ssh { ssh_address } => Some(ssh_address.into()),
        ServerAddress::Localhost => None,
      },
      path: PathBuf::from(path.as_ref()),
    }
  }
}

impl From<&Path> for ScpParam {
  fn from(value: &Path) -> Self {
    Self {
      server: None,
      path: PathBuf::from(value),
    }
  }
}

impl From<&ScpParam> for OsString {
  fn from(value: &ScpParam) -> Self {
    let mut builder = osf!();
    if let Some(server) = &value.server {
      builder += format!("{server}:");
    }
    builder += &value.path;
    builder.build()
  }
}

impl Display for ShellCommand {
  fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    write!(
      f,
      "{}",
      command_to_string(&build_command_from_shell_command(self))
    )
  }
}

fn command_to_string(command: &Command) -> String {
  once(command.get_program().to_string_lossy().to_string())
    .chain(command.get_args().map(|arg| {
      let arg_str = arg.to_string_lossy();
      if arg_str.contains(' ') {
        format!("\"{arg_str}\"")
      } else {
        arg_str.to_string()
      }
    }))
    .collect::<Vec<_>>()
    .join(" ")
}

macro_rules! cmd {
  ($programm: expr $(, $arg:expr )*)=> {{
    #[allow(unused_mut)]
    let mut cmd = Command::new($programm);
    $( cmd.arg($arg); )*
    cmd
  }};
}

pub fn build_command_from_shell_command(shell_command: &ShellCommand) -> Command {
  match shell_command {
    ShellCommand::Ssh {
      address,
      server_command,
    } => cmd!(
      "ssh",
      address,
      match server_command {
        ServerCommand::Realpath { path } => osf!("realpath -e ") + path,
        ServerCommand::Ls { dir } => osf!("ls ") + dir,
        ServerCommand::Rm { file } => osf!("rm ") + file,
        ServerCommand::Mv {
          source,
          destination,
        } => osf!("mv ") + source + " " + destination,
        ServerCommand::Execute {
          working_directory,
          command,
        } => osf!("cd ") + working_directory + "; " + command,
      }
    ),
    ShellCommand::Scp {
      source,
      destination,
    } => cmd!("scp", OsString::from(source), OsString::from(destination)),
    ShellCommand::SshAgent => cmd!("ssh-agent", "-s"),
    ShellCommand::ShhAdd => cmd!("ssh-add"),
    ShellCommand::Editor(args) => {
      let mut args = args.clone();
      let mut cmd = cmd!(args.remove(0));
      cmd.args(args);
      cmd
    }
    ShellCommand::Execute {
      working_directory,
      command,
    } => {
      let mut args = command.clone();
      let mut cmd = cmd!(args.remove(0));
      cmd.args(args).current_dir(working_directory);
      cmd
    }
  }
}

pub trait ShellInterface {
  fn run_command(&mut self, command: ShellCommand) -> CommandResult<ExitStatus, StartError>;
  fn collect_command_output(
    &mut self,
    command: ShellCommand,
  ) -> CommandResult<CommandOutput, StartError>;
}

pub trait MaybeCast<T> {
  fn maybe_cast(&self) -> Option<&T>;
}

impl<T> MaybeCast<T> for T {
  fn maybe_cast(&self) -> Option<&T> {
    Some(self)
  }
}

#[derive(Debug)]
pub enum LoggedRunOutput {
  ExitStatus(ExitStatus),
  CommandOutput(CommandOutput),
}

impl From<ExitStatus> for LoggedRunOutput {
  fn from(value: ExitStatus) -> Self {
    Self::ExitStatus(value)
  }
}

impl From<CommandOutput> for LoggedRunOutput {
  fn from(value: CommandOutput) -> Self {
    Self::CommandOutput(value)
  }
}

impl AsRef<ExitStatus> for LoggedRunOutput {
  fn as_ref(&self) -> &ExitStatus {
    match self {
      LoggedRunOutput::ExitStatus(status) => status,
      LoggedRunOutput::CommandOutput(output) => output.as_ref(),
    }
  }
}

impl MaybeCast<CommandOutput> for LoggedRunOutput {
  fn maybe_cast(&self) -> Option<&CommandOutput> {
    match self {
      LoggedRunOutput::ExitStatus(_) => None,
      LoggedRunOutput::CommandOutput(output) => Some(output),
    }
  }
}

#[derive(Debug)]
pub struct CommandResult<T, E> {
  pub command: ShellCommand,
  pub result: Result<T, E>,
}

impl<T, E> CommandResult<T, E> {
  pub fn into_result(self) -> Result<T, CommandError<E>> {
    self.result.map_err(|error| CommandError {
      command: self.command,
      error,
    })
  }
}

impl<T> CommandResult<T, StartError> {
  pub fn and_expect_success(self) -> CommandResult<T, ExecutionError<T>>
  where
    T: AsRef<ExitStatus>,
  {
    CommandResult {
      result: self.result.map_err(ExecutionError::from).and_then(|t| {
        if t.as_ref().success {
          Ok(t)
        } else {
          Err(ExecutionError::BadExitStatus(t))
        }
      }),
      command: self.command,
    }
  }
}

impl<T> CommandResult<T, ExecutionError<T>> {
  pub fn into_result_with_error_logging(
    self,
    logger: &Logger,
  ) -> Result<T, CommandError<ExecutionError<T>>>
  where
    T: MaybeCast<CommandOutput>,
  {
    self.result.map_err(|error| {
      if let ExecutionError::BadExitStatus(t) = &error {
        if let Some(output) = t.maybe_cast() {
          log!(logger, error, "{}", output.stdout.to_string_lossy());
          log!(logger, error, "{}", output.stderr.to_string_lossy());
        }
      }
      CommandError {
        command: self.command,
        error,
      }
    })
  }
}

#[derive(Debug, Clone)]
pub struct CommandOutput {
  pub stdout: OsString,
  pub stderr: OsString,
  pub status: ExitStatus,
}

impl From<Output> for CommandOutput {
  fn from(value: Output) -> Self {
    Self {
      stdout: os_string_from_ssh_output(value.stdout),
      stderr: os_string_from_ssh_output(value.stderr),
      status: value.status.into(),
    }
  }
}

pub fn os_string_from_ssh_output(output: Vec<u8>) -> OsString {
  #[cfg(unix)]
  {
    use std::os::unix::ffi::OsStringExt;
    OsString::from_vec(output)
  }

  #[cfg(windows)]
  {
    use std::os::windows::ffi::OsStringExt;
    OsString::from_wide(output.iter().map(|&b| b as u16).collect())
  }
}

impl AsRef<ExitStatus> for CommandOutput {
  fn as_ref(&self) -> &ExitStatus {
    &self.status
  }
}

#[derive(Debug, Clone)]
pub struct ExitStatus {
  pub success: bool,
  pub string_form: String,
  #[allow(dead_code)]
  pub code: Option<i32>,
}

impl From<process::ExitStatus> for ExitStatus {
  fn from(value: process::ExitStatus) -> Self {
    Self {
      success: value.success(),
      string_form: value.to_string(),
      code: value.code(),
    }
  }
}

impl AsRef<ExitStatus> for ExitStatus {
  fn as_ref(&self) -> &ExitStatus {
    self
  }
}

impl Display for ExitStatus {
  fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    Display::fmt(&self.string_form, f)
  }
}

#[derive(Debug)]
pub struct CommandError<E> {
  pub command: ShellCommand,
  pub error: E,
}

impl<E> From<CommandError<E>> for String
where
  E: Display,
{
  fn from(value: CommandError<E>) -> Self {
    value.to_string()
  }
}

impl<E> Display for CommandError<E>
where
  E: Display,
{
  fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    write!(
      f,
      "Error while running command '{}': {}",
      self.command, self.error
    )
  }
}

impl<E> Error for CommandError<E> where E: Error {}

#[derive(Debug)]
pub struct StartError(io::Error);

impl From<io::Error> for StartError {
  fn from(value: io::Error) -> Self {
    StartError(value)
  }
}

impl Display for StartError {
  fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    write!(f, "Failed to run command: {}", self.0)
  }
}

impl Error for StartError {}

#[derive(Debug)]
pub enum ExecutionError<T> {
  StartError(StartError),
  BadExitStatus(T),
}

impl<T> From<StartError> for ExecutionError<T> {
  fn from(value: StartError) -> Self {
    Self::StartError(value)
  }
}

impl<T> Display for ExecutionError<T>
where
  T: AsRef<ExitStatus>,
{
  fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
    match self {
      ExecutionError::StartError(e) => Display::fmt(e, f),
      ExecutionError::BadExitStatus(status) => {
        write!(f, "execution failed with {}", status.as_ref())
      }
    }
  }
}

impl<T> Error for ExecutionError<T> where T: AsRef<ExitStatus> + Debug {}

#[cfg(test)]
mod test_commands {
  use crate::shell_interface::{ScpParam, ServerCommand, ShellCommand};
  use std::path::PathBuf;

  #[test]
  fn test_to_string() {
    assert_eq!(
      ShellCommand::Ssh {
        address: "crea".to_string(),
        server_command: ServerCommand::Realpath {
          path: PathBuf::from("plugins/*.jar")
        }
      }
      .to_string(),
      r#"ssh crea "realpath -e plugins/*.jar""#
    );
    assert_eq!(
      ShellCommand::Ssh {
        address: "crea".to_string(),
        server_command: ServerCommand::Ls {
          dir: PathBuf::from("creative/plugins")
        }
      }
      .to_string(),
      r#"ssh crea "ls creative/plugins""#
    );
    assert_eq!(
      ShellCommand::Ssh {
        address: "crea".to_string(),
        server_command: ServerCommand::Rm {
          file: PathBuf::from("foo.txt")
        },
      }
      .to_string(),
      r#"ssh crea "rm foo.txt""#
    );
    assert_eq!(
      ShellCommand::Ssh {
        address: "crea".to_string(),
        server_command: ServerCommand::Mv {
          source: PathBuf::from("foo"),
          destination: PathBuf::from("bar")
        }
      }
      .to_string(),
      r#"ssh crea "mv foo bar""#
    );
    assert_eq!(
      ShellCommand::Ssh {
        address: "crea".to_string(),
        server_command: ServerCommand::Execute {
          working_directory: PathBuf::from(".."),
          command: "sudo rm -rf *".into(),
        }
      }
      .to_string(),
      r#"ssh crea "cd ..; sudo rm -rf *""#
    );
    assert_eq!(
      ShellCommand::Scp {
        source: ScpParam {
          server: None,
          path: PathBuf::from("target/mssh")
        },
        destination: ScpParam {
          server: Some("crea".into()),
          path: PathBuf::from("/usr/bin")
        },
      }
      .to_string(),
      r#"scp target/mssh crea:/usr/bin"#
    );
    assert_eq!(ShellCommand::SshAgent.to_string(), r#"ssh-agent -s"#);
    assert_eq!(ShellCommand::ShhAdd.to_string(), r#"ssh-add"#);
    assert_eq!(
      ShellCommand::Editor(vec!["kate".into(), "-b".into(), "test.txt".into()]).to_string(),
      r#"kate -b test.txt"#
    );
    assert_eq!(
      ShellCommand::Execute {
        working_directory: PathBuf::from("/home/me/server"),
        command: vec!["java".into(), "-jar".into(), "paper.jar".into()]
      }
      .to_string(),
      r#"java -jar paper.jar"#
    );
  }
}