use crate::commands::ssh_helpers::{prompt_for_input, prompt_for_password};
use crate::config::SshConfig;
use async_ssh2_tokio::client::{AuthMethod, Client, ServerCheckMethod};
use russh::{ChannelMsg, Sig};
use tokio::io::{stderr, stdout, AsyncWriteExt};
use tokio::signal;
use tracing::debug;
pub async fn run_remote_logs(
project: Option<String>,
ssh_host: Option<String>,
ssh_username: Option<String>,
ssh_password: Option<String>,
debug_mode: bool,
) -> Result<(), String> {
let mut config = SshConfig::load().map_err(|e| format!("Failed to load SSH config: {}", e))?;
let mut config_dirty = false;
let resolved_host = resolve_or_prompt(
ssh_host,
&mut config.host,
"Enter SSH host: ",
&mut config_dirty,
)?;
let resolved_username = resolve_or_prompt(
ssh_username,
&mut config.username,
"Enter SSH username: ",
&mut config_dirty,
)?;
let resolved_password = resolve_or_prompt_password(
ssh_password,
&mut config.password,
"Enter SSH password: ",
&mut config_dirty,
)?;
if config_dirty {
config
.save()
.map_err(|e| format!("Failed to save SSH config: {}", e))?;
}
let command = build_remote_command(project.as_deref());
let wrapped_command = wrap_command_with_login_shell(&command);
if debug_mode {
debug!(
"SSH logs target => host: {}, user: {}, command: {}",
resolved_host, resolved_username, wrapped_command
);
}
let auth = AuthMethod::with_password(&resolved_password);
let client = Client::connect(
(resolved_host.as_str(), 22),
resolved_username.as_str(),
auth,
ServerCheckMethod::NoCheck,
)
.await
.map_err(|e| format!("SSH connection failed: {}", e))?;
let streaming_result = stream_remote_logs(&client, &wrapped_command).await;
if let Err(err) = client.disconnect().await {
debug!("Failed to cleanly disconnect SSH session: {}", err);
}
streaming_result
}
fn resolve_or_prompt(
provided: Option<String>,
stored: &mut Option<String>,
prompt: &str,
dirty: &mut bool,
) -> Result<String, String> {
if let Some(value) = provided {
if stored.as_ref() != Some(&value) {
*stored = Some(value.clone());
*dirty = true;
}
return Ok(value);
}
if let Some(value) = stored.clone() {
return Ok(value);
}
let value = prompt_for_input(prompt)?;
*stored = Some(value.clone());
*dirty = true;
Ok(value)
}
fn resolve_or_prompt_password(
provided: Option<String>,
stored: &mut Option<String>,
prompt: &str,
dirty: &mut bool,
) -> Result<String, String> {
if let Some(value) = provided {
if stored.as_ref() != Some(&value) {
*stored = Some(value.clone());
*dirty = true;
}
return Ok(value);
}
if let Some(value) = stored.clone() {
return Ok(value);
}
let value = prompt_for_password(prompt)?;
*stored = Some(value.clone());
*dirty = true;
Ok(value)
}
fn build_remote_command(project: Option<&str>) -> String {
if let Some(project_name) = project {
format!("xbp logs {}", shell_quote(project_name))
} else {
"xbp logs".to_string()
}
}
fn wrap_command_with_login_shell(command: &str) -> String {
format!("bash -lc {}", shell_quote(command))
}
fn shell_quote(value: &str) -> String {
if value.is_empty() {
return "''".to_string();
}
if value
.chars()
.all(|ch| ch.is_ascii_alphanumeric() || "-_./:@".contains(ch))
{
return value.to_string();
}
format!("'{}'", value.replace('\'', "'\\''"))
}
async fn stream_remote_logs(client: &Client, command: &str) -> Result<(), String> {
let mut channel = client
.get_channel()
.await
.map_err(|e| format!("Failed to open SSH channel: {}", e))?;
channel
.exec(true, command)
.await
.map_err(|e| format!("Failed to execute remote command: {}", e))?;
let mut stdout = stdout();
let mut stderr = stderr();
let mut exit_status: Option<u32> = None;
let mut ctrl_c_pressed = false;
let ctrl_c = signal::ctrl_c();
tokio::pin!(ctrl_c);
loop {
tokio::select! {
msg = channel.wait() => match msg {
Some(ChannelMsg::Data { ref data }) => {
stdout
.write_all(data)
.await
.map_err(|e| format!("Failed to write remote stdout: {}", e))?;
stdout
.flush()
.await
.map_err(|e| format!("Failed to flush stdout: {}", e))?;
}
Some(ChannelMsg::ExtendedData { ref data, ext }) => {
if ext == 1 {
stderr
.write_all(data)
.await
.map_err(|e| format!("Failed to write remote stderr: {}", e))?;
stderr
.flush()
.await
.map_err(|e| format!("Failed to flush stderr: {}", e))?;
}
}
Some(ChannelMsg::ExitStatus { exit_status: status }) => {
exit_status = Some(status);
}
Some(ChannelMsg::ExitSignal { signal_name, .. }) => {
if exit_status.is_none() {
exit_status = Some(signal_to_exit_status(&signal_name));
}
}
Some(_) => {}
None => break,
},
_ = &mut ctrl_c => {
ctrl_c_pressed = true;
if let Err(err) = channel.signal(Sig::INT).await {
debug!("Failed to send interrupt signal: {}", err);
}
}
}
}
if ctrl_c_pressed {
return Ok(());
}
if let Some(status) = exit_status {
if status == 0 {
Ok(())
} else {
Err(format!("Remote command exited with status: {}", status))
}
} else {
Err("Remote command closed without exit status".to_string())
}
}
fn signal_to_exit_status(signal: &Sig) -> u32 {
match signal {
Sig::INT => 130,
Sig::TERM => 143,
Sig::QUIT => 131,
Sig::KILL => 137,
Sig::HUP => 129,
Sig::PIPE => 141,
_ => 128,
}
}