johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::process::Stdio;
use std::time::Duration;

use tokio::io::AsyncWriteExt;
use tokio::process::Command;

const COMMAND_TIMEOUT: Duration = Duration::from_secs(60);

#[derive(Clone, Debug)]
pub struct Runner {
    ssh: Option<String>,
    sudo: bool,
    kubectl: String,
}

impl Runner {
    pub fn new(ssh: Option<String>, sudo: bool, kubectl: Option<String>) -> Self {
        Self {
            ssh,
            sudo,
            kubectl: kubectl.unwrap_or_else(|| "kubectl".to_string()),
        }
    }

    pub fn kubectl_path(&self) -> &str {
        &self.kubectl
    }

    pub fn ssh_host(&self) -> Option<&str> {
        self.ssh.as_deref()
    }

    pub async fn kubectl(&self, args: &[String], stdin: Option<&str>) -> Result<String, String> {
        let mut command = Vec::with_capacity(args.len() + 1);
        command.push(self.kubectl.clone());
        command.extend(args.iter().cloned());
        self.run(&command, stdin).await
    }

    pub async fn shell(&self, script: &str) -> Result<String, String> {
        self.shell_with_stdin(script, None).await
    }

    pub async fn shell_with_stdin(
        &self,
        script: &str,
        stdin: Option<&str>,
    ) -> Result<String, String> {
        let command = vec!["sh".to_string(), "-c".to_string(), script.to_string()];
        self.run(&command, stdin).await
    }

    /// Spawns `kubectl <args>` with stdout piped for streaming (watches).
    /// The child is killed when dropped.
    pub fn spawn_kubectl_stream(&self, args: &[String]) -> Result<tokio::process::Child, String> {
        let mut argv = Vec::with_capacity(args.len() + 1);
        argv.push(self.kubectl.clone());
        argv.extend(args.iter().cloned());
        let mut command = self.command(&argv);
        command
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .kill_on_drop(true);
        command
            .spawn()
            .map_err(|error| format!("failed to spawn {:?}: {error}", argv[0]))
    }

    /// The local or `ssh` command for `args`, with `sudo -n` when enabled.
    fn command(&self, args: &[String]) -> Command {
        let command = match &self.ssh {
            Some(host) => {
                let mut remote = Vec::with_capacity(args.len() + 2);
                if self.sudo {
                    remote.push("sudo".to_string());
                    remote.push("-n".to_string());
                }
                remote.extend(args.iter().cloned());
                let script = remote
                    .iter()
                    .map(|argument| shell_quote(argument))
                    .collect::<Vec<_>>()
                    .join(" ");
                let mut command = Command::new("ssh");
                command
                    .arg("-o")
                    .arg("BatchMode=yes")
                    .arg("-o")
                    .arg("ConnectTimeout=10")
                    .arg(host)
                    .arg(script);
                command
            }
            None => {
                let (program, rest) = args.split_first().expect("checked non-empty");
                let mut command = if self.sudo {
                    let mut command = Command::new("sudo");
                    command.arg("-n").arg(program);
                    command
                } else {
                    Command::new(program)
                };
                command.args(rest);
                command
            }
        };
        command
    }

    pub async fn run(&self, args: &[String], stdin: Option<&str>) -> Result<String, String> {
        if args.is_empty() {
            return Err("refusing to run an empty command".to_string());
        }

        let mut command = self.command(args);
        command
            .stdin(if stdin.is_some() {
                Stdio::piped()
            } else {
                Stdio::null()
            })
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());

        let mut child = command
            .spawn()
            .map_err(|error| format!("failed to spawn {:?}: {error}", args[0]))?;

        if let Some(input) = stdin {
            let mut pipe = child
                .stdin
                .take()
                .ok_or_else(|| "failed to open command stdin".to_string())?;
            pipe.write_all(input.as_bytes())
                .await
                .map_err(|error| format!("failed to write command stdin: {error}"))?;
            pipe.shutdown()
                .await
                .map_err(|error| format!("failed to close command stdin: {error}"))?;
        }

        let output = tokio::time::timeout(COMMAND_TIMEOUT, child.wait_with_output())
            .await
            .map_err(|_| {
                format!(
                    "command timed out after {}s: {}",
                    COMMAND_TIMEOUT.as_secs(),
                    args.join(" ")
                )
            })?
            .map_err(|error| format!("command failed: {error}"))?;

        let stdout = String::from_utf8_lossy(&output.stdout).to_string();
        if output.status.success() {
            return Ok(stdout);
        }

        let stderr = String::from_utf8_lossy(&output.stderr);
        let detail = if stderr.trim().is_empty() {
            stdout
        } else {
            stderr.to_string()
        };
        Err(format!(
            "command exited with {}: {}",
            output.status,
            detail.trim()
        ))
    }
}

pub fn shell_quote(value: &str) -> String {
    if value
        .chars()
        .all(|character| character.is_ascii_alphanumeric() || "-_./:=@".contains(character))
    {
        return value.to_string();
    }
    format!("'{}'", value.replace('\'', "'\\''"))
}

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

    #[test]
    fn quotes_only_what_the_remote_shell_needs() {
        assert_eq!(shell_quote("kubectl"), "kubectl");
        assert_eq!(
            shell_quote("/var/lib/kcp/admin.kubeconfig"),
            "/var/lib/kcp/admin.kubeconfig"
        );
        assert_eq!(
            shell_quote("*kcp/admin.kubeconfig"),
            "'*kcp/admin.kubeconfig'"
        );
        assert_eq!(shell_quote("a b"), "'a b'");
        assert_eq!(shell_quote("it's"), "'it'\\''s'");
    }
}