harmony/harmony-k8s/src/pod.rs

use std::time::Duration;

use k8s_openapi::api::core::v1::Pod;
use kube::{
    Error,
    api::{Api, AttachParams, ListParams},
    error::DiscoveryError,
    runtime::reflector::Lookup,
};
use log::debug;
use tokio::io::AsyncReadExt;
use tokio::time::sleep;

use crate::client::K8sClient;

/// Result of executing a command in a pod. Carries both streams plus the
/// reported status (`"Success"` or `"Failure"`) so callers can react
/// without losing output — important when the executed command writes
/// structured payloads to stdout on non-zero exits (the Vault / OpenBao
/// CLIs do this with `-format=json`).
#[derive(Debug, Clone)]
pub struct ExecOutput {
    pub status: String,
    pub stdout: String,
    pub stderr: String,
}

impl ExecOutput {
    pub fn is_success(&self) -> bool {
        self.status == "Success"
    }
}

impl K8sClient {
    pub async fn get_pod(&self, name: &str, namespace: Option<&str>) -> Result<Option<Pod>, Error> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };
        api.get_opt(name).await
    }

    pub async fn wait_for_pod_ready(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
    ) -> Result<(), Error> {
        let mut elapsed = 0u64;
        let interval = 5u64;
        let timeout_secs = 120u64;
        loop {
            if let Some(p) = self.get_pod(pod_name, namespace).await?
                && let Some(phase) = p.status.and_then(|s| s.phase)
                && phase.to_lowercase() == "running"
            {
                return Ok(());
            }
            if elapsed >= timeout_secs {
                return Err(Error::Discovery(DiscoveryError::MissingResource(format!(
                    "Pod '{}' in '{}' did not become ready within {timeout_secs}s",
                    pod_name,
                    namespace.unwrap_or("<default>"),
                ))));
            }
            sleep(Duration::from_secs(interval)).await;
            elapsed += interval;
        }
    }

    /// Polls a pod until it reaches `Succeeded` or `Failed`, then returns its
    /// logs.  Used internally by node operations.
    pub(crate) async fn wait_for_pod_completion(
        &self,
        name: &str,
        namespace: &str,
    ) -> Result<String, Error> {
        let api: Api<Pod> = Api::namespaced(self.client.clone(), namespace);
        let poll_interval = Duration::from_secs(2);
        for _ in 0..60 {
            sleep(poll_interval).await;
            let p = api.get(name).await?;
            match p.status.and_then(|s| s.phase).as_deref() {
                Some("Succeeded") => {
                    let logs = api
                        .logs(name, &Default::default())
                        .await
                        .unwrap_or_default();
                    debug!("Pod {namespace}/{name} succeeded. Logs: {logs}");
                    return Ok(logs);
                }
                Some("Failed") => {
                    let logs = api
                        .logs(name, &Default::default())
                        .await
                        .unwrap_or_default();
                    debug!("Pod {namespace}/{name} failed. Logs: {logs}");
                    return Err(Error::Discovery(DiscoveryError::MissingResource(format!(
                        "Pod '{name}' failed.\n{logs}"
                    ))));
                }
                _ => {}
            }
        }
        Err(Error::Discovery(DiscoveryError::MissingResource(format!(
            "Timed out waiting for pod '{name}'"
        ))))
    }

    /// Execute a command in the first pod matching `{label}={name}`.
    pub async fn exec_app_capture_output(
        &self,
        name: String,
        label: String,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<String, String> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };
        let pod_list = api
            .list(&ListParams::default().labels(&format!("{label}={name}")))
            .await
            .expect("Failed to list pods");

        let pod_name = pod_list
            .items
            .first()
            .expect("No matching pod")
            .name()
            .expect("Pod has no name")
            .into_owned();

        match api
            .exec(
                &pod_name,
                command,
                &AttachParams::default().stdout(true).stderr(true),
            )
            .await
        {
            Err(e) => Err(e.to_string()),
            Ok(mut process) => {
                let status = process
                    .take_status()
                    .expect("No status handle")
                    .await
                    .expect("Status channel closed");

                if let Some(s) = status.status {
                    let mut buf = String::new();
                    if let Some(mut stdout) = process.stdout() {
                        stdout
                            .read_to_string(&mut buf)
                            .await
                            .map_err(|e| format!("Failed to read stdout: {e}"))?;
                    }
                    debug!("exec status: {} - {:?}", s, status.details);
                    if s == "Success" { Ok(buf) } else { Err(s) }
                } else {
                    Err("No inner status from pod exec".to_string())
                }
            }
        }
    }

    /// Execute a command in the first pod matching
    /// `app.kubernetes.io/name={name}`.
    pub async fn exec_app(
        &self,
        name: String,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<(), String> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };
        let pod_list = api
            .list(&ListParams::default().labels(&format!("app.kubernetes.io/name={name}")))
            .await
            .expect("Failed to list pods");

        let pod_name = pod_list
            .items
            .first()
            .expect("No matching pod")
            .name()
            .expect("Pod has no name")
            .into_owned();

        match api.exec(&pod_name, command, &AttachParams::default()).await {
            Err(e) => Err(e.to_string()),
            Ok(mut process) => {
                let status = process
                    .take_status()
                    .expect("No status handle")
                    .await
                    .expect("Status channel closed");

                if let Some(s) = status.status {
                    debug!("exec status: {} - {:?}", s, status.details);
                    if s == "Success" { Ok(()) } else { Err(s) }
                } else {
                    Err("No inner status from pod exec".to_string())
                }
            }
        }
    }

    /// Execute a command in a specific pod by name, capturing stdout.
    ///
    /// Returns the captured stdout on success. On failure, the error string
    /// includes stderr output from the remote command.
    pub async fn exec_pod_capture_output(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<String, String> {
        // Convenience wrapper — preserved for backward compatibility:
        // returns stdout on success, stderr on non-zero exit, and drops
        // the other stream on each path. New callers should prefer
        // `exec_pod_capture` which exposes both streams unconditionally,
        // important when the failed command writes structured output to
        // stdout (e.g. `bao status -format=json` exits 2 on a sealed or
        // uninitialised vault but still emits its JSON payload).
        let output = self.exec_pod_capture(pod_name, namespace, command).await?;
        if output.is_success() {
            Ok(output.stdout)
        } else {
            Err(output.stderr)
        }
    }

    /// Execute a command in a pod and capture both stdout and stderr
    /// regardless of the process exit status. The outer `Result` is
    /// reserved for transport / API failures (couldn't reach the pod,
    /// stream read error). A non-zero exit is reflected in
    /// `ExecOutput::status` and the caller decides how to react.
    pub async fn exec_pod_capture(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<ExecOutput, String> {
        let api: Api<Pod> = match namespace {
            Some(ns) => Api::namespaced(self.client.clone(), ns),
            None => Api::default_namespaced(self.client.clone()),
        };

        let mut process = api
            .exec(
                pod_name,
                command,
                &AttachParams::default().stdout(true).stderr(true),
            )
            .await
            .map_err(|e| e.to_string())?;

        let status = process
            .take_status()
            .expect("No status handle")
            .await
            .expect("Status channel closed");

        let mut stdout = String::new();
        if let Some(mut s) = process.stdout() {
            s.read_to_string(&mut stdout)
                .await
                .map_err(|e| format!("Failed to read stdout: {e}"))?;
        }

        let mut stderr = String::new();
        if let Some(mut s) = process.stderr() {
            s.read_to_string(&mut stderr)
                .await
                .map_err(|e| format!("Failed to read stderr: {e}"))?;
        }

        let status_field = status
            .status
            .ok_or_else(|| "No inner status from pod exec".to_string())?;
        debug!("exec_pod status: {} - {:?}", status_field, status.details);

        Ok(ExecOutput {
            status: status_field,
            stdout,
            stderr,
        })
    }

    /// Execute a command in a specific pod by name (no output capture).
    pub async fn exec_pod(
        &self,
        pod_name: &str,
        namespace: Option<&str>,
        command: Vec<&str>,
    ) -> Result<(), String> {
        self.exec_pod_capture_output(pod_name, namespace, command)
            .await
            .map(|_| ())
    }
}