harmony/harmony-k8s/src/helper.rs

use std::collections::BTreeMap;
use std::time::Duration;

use crate::KubernetesDistribution;

use super::bundle::ResourceBundle;
use super::config::PRIVILEGED_POD_IMAGE;
use k8s_openapi::api::core::v1::{
    Container, HostPathVolumeSource, Pod, PodSpec, SecurityContext, Volume, VolumeMount,
};
use k8s_openapi::api::rbac::v1::{ClusterRoleBinding, RoleRef, Subject};
use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
use kube::api::DynamicObject;
use kube::error::DiscoveryError;
use log::{debug, error, info, warn};
use serde::de::DeserializeOwned;

#[derive(Debug)]
pub struct PrivilegedPodConfig {
    pub name: String,
    pub namespace: String,
    pub node_name: String,
    pub container_name: String,
    pub command: Vec<String>,
    pub volumes: Vec<Volume>,
    pub volume_mounts: Vec<VolumeMount>,
    pub host_pid: bool,
    pub host_network: bool,
}

impl Default for PrivilegedPodConfig {
    fn default() -> Self {
        Self {
            name: "privileged-pod".to_string(),
            namespace: "harmony".to_string(),
            node_name: "".to_string(),
            container_name: "privileged-container".to_string(),
            command: vec![],
            volumes: vec![],
            volume_mounts: vec![],
            host_pid: false,
            host_network: false,
        }
    }
}

pub fn build_privileged_pod(
    config: PrivilegedPodConfig,
    k8s_distribution: &KubernetesDistribution,
) -> Pod {
    let annotations = match k8s_distribution {
        KubernetesDistribution::OpenshiftFamily => Some(BTreeMap::from([
            ("openshift.io/scc".to_string(), "privileged".to_string()),
            (
                "openshift.io/required-scc".to_string(),
                "privileged".to_string(),
            ),
        ])),
        _ => None,
    };

    Pod {
        metadata: ObjectMeta {
            name: Some(config.name),
            namespace: Some(config.namespace),
            annotations,
            ..Default::default()
        },
        spec: Some(PodSpec {
            node_name: Some(config.node_name),
            restart_policy: Some("Never".to_string()),
            host_pid: Some(config.host_pid),
            host_network: Some(config.host_network),
            containers: vec![Container {
                name: config.container_name,
                image: Some(PRIVILEGED_POD_IMAGE.to_string()),
                command: Some(config.command),
                security_context: Some(SecurityContext {
                    privileged: Some(true),
                    ..Default::default()
                }),
                volume_mounts: Some(config.volume_mounts),
                ..Default::default()
            }],
            volumes: Some(config.volumes),
            ..Default::default()
        }),
        ..Default::default()
    }
}

pub fn host_root_volume() -> (Volume, VolumeMount) {
    (
        Volume {
            name: "host".to_string(),
            host_path: Some(HostPathVolumeSource {
                path: "/".to_string(),
                ..Default::default()
            }),
            ..Default::default()
        },
        VolumeMount {
            name: "host".to_string(),
            mount_path: "/host".to_string(),
            ..Default::default()
        },
    )
}

/// Build a ResourceBundle containing a privileged pod and any required RBAC.
///
/// This function implements the Resource Bundle pattern to encapsulate platform-specific
/// security requirements for running privileged operations on nodes.
///
/// # Platform-Specific Behavior
///
/// - **OpenShift**: Creates a ClusterRoleBinding to grant the default ServiceAccount
///   access to the `system:openshift:scc:privileged` ClusterRole, which allows the pod
///   to use the privileged Security Context Constraint (SCC).
/// - **Standard Kubernetes/K3s**: Only creates the Pod resource, as these distributions
///   use standard PodSecurityPolicy or don't enforce additional security constraints.
///
/// # Arguments
///
/// * `config` - Configuration for the privileged pod (name, namespace, command, etc.)
/// * `k8s_distribution` - The detected Kubernetes distribution to determine RBAC requirements
///
/// # Returns
///
/// A `ResourceBundle` containing 1-2 resources:
/// - ClusterRoleBinding (OpenShift only)
/// - Pod (all distributions)
///
/// # Example
///
/// ```
/// use harmony_k8s::helper::{build_privileged_bundle, PrivilegedPodConfig};
/// use harmony_k8s::KubernetesDistribution;
/// let bundle = build_privileged_bundle(
///     PrivilegedPodConfig {
///         name: "network-setup".to_string(),
///         namespace: "default".to_string(),
///         node_name: "worker-01".to_string(),
///         container_name: "setup".to_string(),
///         command: vec!["nmcli".to_string(), "connection".to_string(), "reload".to_string()],
///         ..Default::default()
///     },
///     &KubernetesDistribution::OpenshiftFamily,
/// );
/// // Bundle now contains ClusterRoleBinding + Pod
/// ```
pub fn build_privileged_bundle(
    config: PrivilegedPodConfig,
    k8s_distribution: &KubernetesDistribution,
) -> ResourceBundle {
    debug!(
        "Building privileged bundle for config {config:#?} on distribution {k8s_distribution:?}"
    );
    let mut bundle = ResourceBundle::new();
    let pod_name = config.name.clone();
    let namespace = config.namespace.clone();

    // 1. On OpenShift, create RBAC binding to privileged SCC
    if let KubernetesDistribution::OpenshiftFamily = k8s_distribution {
        // The default ServiceAccount needs to be bound to the privileged SCC
        // via the system:openshift:scc:privileged ClusterRole
        let crb = ClusterRoleBinding {
            metadata: ObjectMeta {
                name: Some(format!("{}-scc-binding", pod_name)),
                ..Default::default()
            },
            role_ref: RoleRef {
                api_group: "rbac.authorization.k8s.io".to_string(),
                kind: "ClusterRole".to_string(),
                name: "system:openshift:scc:privileged".to_string(),
            },
            subjects: Some(vec![Subject {
                kind: "ServiceAccount".to_string(),
                name: "default".to_string(),
                namespace: Some(namespace.clone()),
                api_group: None,
            }]),
        };
        bundle.add(crb);
    }

    // 2. Build the privileged pod
    let pod = build_privileged_pod(config, k8s_distribution);
    bundle.add(pod);

    bundle
}

/// Action to take when a drain operation times out.
pub enum DrainTimeoutAction {
    /// Accept the partial drain and continue
    Accept,
    /// Retry the drain for another timeout period
    Retry,
    /// Abort the drain operation
    Abort,
}

/// Prompts the user to confirm acceptance of a partial drain.
///
/// Returns `Ok(true)` if the user confirms acceptance, `Ok(false)` if the user
/// chooses to retry or abort, and `Err` if the prompt system fails entirely.
pub fn prompt_drain_timeout_action(
    node_name: &str,
    pending_count: usize,
    timeout_duration: Duration,
) -> Result<DrainTimeoutAction, kube::Error> {
    let prompt_msg = format!(
        "Drain operation timed out on node '{}' with {} pod(s) remaining. What would you like to do?",
        node_name, pending_count
    );

    loop {
        let choices = vec![
            "Accept drain failure (requires confirmation)".to_string(),
            format!("Retry drain for another {:?}", timeout_duration),
            "Abort operation".to_string(),
        ];

        let selection = inquire::Select::new(&prompt_msg, choices)
            .with_help_message("Use arrow keys to navigate, Enter to select")
            .prompt()
            .map_err(|e| {
                kube::Error::Discovery(DiscoveryError::MissingResource(format!(
                    "Prompt failed: {}",
                    e
                )))
            })?;

        if selection.starts_with("Accept") {
            // Require typed confirmation - retry until correct or user cancels
            let required_confirmation = format!("yes-accept-drain:{}={}", node_name, pending_count);

            let confirmation_prompt = format!(
                "To accept this partial drain, type exactly: {}",
                required_confirmation
            );

            match inquire::Text::new(&confirmation_prompt)
                .with_help_message(&format!(
                    "This action acknowledges {} pods will remain on the node",
                    pending_count
                ))
                .prompt()
            {
                Ok(input) if input == required_confirmation => {
                    warn!(
                        "User accepted partial drain of node '{}' with {} pods remaining (confirmation: {})",
                        node_name, pending_count, required_confirmation
                    );
                    return Ok(DrainTimeoutAction::Accept);
                }
                Ok(input) => {
                    warn!(
                        "Confirmation failed. Expected '{}', got '{}'. Please try again.",
                        required_confirmation, input
                    );
                }
                Err(e) => {
                    // User cancelled (Ctrl+C) or prompt system failed
                    error!("Confirmation prompt cancelled or failed: {}", e);
                    return Ok(DrainTimeoutAction::Abort);
                }
            }
        } else if selection.starts_with("Retry") {
            info!(
                "User chose to retry drain operation for another {:?}",
                timeout_duration
            );
            return Ok(DrainTimeoutAction::Retry);
        } else {
            error!("Drain operation aborted by user");
            return Ok(DrainTimeoutAction::Abort);
        }
    }
}

/// JSON round-trip: DynamicObject → K
///
/// Safe because the DynamicObject was produced by the apiserver from a
/// payload that was originally serialized from K, so the schema is identical.
pub(crate) fn dyn_to_typed<K: DeserializeOwned>(obj: DynamicObject) -> Result<K, kube::Error> {
    serde_json::to_value(obj)
        .and_then(serde_json::from_value)
        .map_err(kube::Error::SerdeError)
}

#[cfg(test)]
mod tests {
    use super::*;
    use pretty_assertions::assert_eq;

    #[test]
    fn test_host_root_volume() {
        let (volume, mount) = host_root_volume();

        assert_eq!(volume.name, "host");
        assert_eq!(volume.host_path.as_ref().unwrap().path, "/");

        assert_eq!(mount.name, "host");
        assert_eq!(mount.mount_path, "/host");
    }

    #[test]
    fn test_build_privileged_pod_minimal() {
        let pod = build_privileged_pod(
            PrivilegedPodConfig {
                name: "minimal-pod".to_string(),
                namespace: "kube-system".to_string(),
                node_name: "node-123".to_string(),
                container_name: "debug-container".to_string(),
                command: vec!["sleep".to_string(), "3600".to_string()],
                ..Default::default()
            },
            &KubernetesDistribution::Default,
        );

        assert_eq!(pod.metadata.name, Some("minimal-pod".to_string()));
        assert_eq!(pod.metadata.namespace, Some("kube-system".to_string()));

        let spec = pod.spec.as_ref().expect("Pod spec should be present");
        assert_eq!(spec.node_name, Some("node-123".to_string()));
        assert_eq!(spec.restart_policy, Some("Never".to_string()));
        assert_eq!(spec.host_pid, Some(false));
        assert_eq!(spec.host_network, Some(false));

        assert_eq!(spec.containers.len(), 1);
        let container = &spec.containers[0];
        assert_eq!(container.name, "debug-container");
        assert_eq!(container.image, Some(PRIVILEGED_POD_IMAGE.to_string()));
        assert_eq!(
            container.command,
            Some(vec!["sleep".to_string(), "3600".to_string()])
        );

        // Security context check
        let sec_ctx = container
            .security_context
            .as_ref()
            .expect("Security context missing");
        assert_eq!(sec_ctx.privileged, Some(true));
    }

    #[test]
    fn test_build_privileged_pod_with_volumes_and_host_access() {
        let (host_vol, host_mount) = host_root_volume();

        let pod = build_privileged_pod(
            PrivilegedPodConfig {
                name: "full-pod".to_string(),
                namespace: "default".to_string(),
                node_name: "node-1".to_string(),
                container_name: "runner".to_string(),
                command: vec!["/bin/sh".to_string()],
                volumes: vec![host_vol.clone()],
                volume_mounts: vec![host_mount.clone()],
                host_pid: true,
                host_network: true,
            },
            &KubernetesDistribution::Default,
        );

        let spec = pod.spec.as_ref().expect("Pod spec should be present");
        assert_eq!(spec.host_pid, Some(true));
        assert_eq!(spec.host_network, Some(true));

        // Check volumes in Spec
        let volumes = spec.volumes.as_ref().expect("Volumes should be present");
        assert_eq!(volumes.len(), 1);
        assert_eq!(volumes[0].name, "host");

        // Check mounts in Container
        let container = &spec.containers[0];
        let mounts = container
            .volume_mounts
            .as_ref()
            .expect("Mounts should be present");
        assert_eq!(mounts.len(), 1);
        assert_eq!(mounts[0].name, "host");
        assert_eq!(mounts[0].mount_path, "/host");
    }

    #[test]
    fn test_build_privileged_pod_structure_correctness() {
        // This test validates that the construction logic puts things in the right places
        // effectively validating the "template".

        let custom_vol = Volume {
            name: "custom-vol".to_string(),
            ..Default::default()
        };
        let custom_mount = VolumeMount {
            name: "custom-vol".to_string(),
            mount_path: "/custom".to_string(),
            ..Default::default()
        };

        let pod = build_privileged_pod(
            PrivilegedPodConfig {
                name: "structure-test".to_string(),
                namespace: "test-ns".to_string(),
                node_name: "test-node".to_string(),
                container_name: "test-container".to_string(),
                command: vec!["cmd".to_string()],
                volumes: vec![custom_vol],
                volume_mounts: vec![custom_mount],
                ..Default::default()
            },
            &KubernetesDistribution::Default,
        );

        // Validate structure depth
        let spec = pod.spec.as_ref().unwrap();

        // 1. Spec level fields
        assert!(spec.node_name.is_some());
        assert!(spec.volumes.is_some());

        // 2. Container level fields
        let container = &spec.containers[0];
        assert!(container.security_context.is_some());
        assert!(container.volume_mounts.is_some());

        // 3. Nested fields
        assert!(
            container
                .security_context
                .as_ref()
                .unwrap()
                .privileged
                .unwrap()
        );
        assert_eq!(spec.volumes.as_ref().unwrap()[0].name, "custom-vol");
        assert_eq!(
            container.volume_mounts.as_ref().unwrap()[0].mount_path,
            "/custom"
        );
    }

    #[test]
    fn test_build_privileged_bundle_default_distribution() {
        let bundle = build_privileged_bundle(
            PrivilegedPodConfig {
                name: "test-bundle".to_string(),
                namespace: "test-ns".to_string(),
                node_name: "node-1".to_string(),
                container_name: "test-container".to_string(),
                command: vec!["echo".to_string(), "hello".to_string()],
                ..Default::default()
            },
            &KubernetesDistribution::Default,
        );

        // For Default distribution, only the Pod should be in the bundle
        assert_eq!(bundle.resources.len(), 1);

        let pod_obj = &bundle.resources[0];
        assert_eq!(pod_obj.metadata.name.as_deref(), Some("test-bundle"));
        assert_eq!(pod_obj.metadata.namespace.as_deref(), Some("test-ns"));
    }

    #[test]
    fn test_build_privileged_bundle_openshift_distribution() {
        let bundle = build_privileged_bundle(
            PrivilegedPodConfig {
                name: "test-bundle-ocp".to_string(),
                namespace: "test-ns".to_string(),
                node_name: "node-1".to_string(),
                container_name: "test-container".to_string(),
                command: vec!["echo".to_string(), "hello".to_string()],
                ..Default::default()
            },
            &KubernetesDistribution::OpenshiftFamily,
        );

        // For OpenShift, both ClusterRoleBinding and Pod should be in the bundle
        assert_eq!(bundle.resources.len(), 2);

        // First resource should be the ClusterRoleBinding
        let crb_obj = &bundle.resources[0];
        assert_eq!(
            crb_obj.metadata.name.as_deref(),
            Some("test-bundle-ocp-scc-binding")
        );

        // Verify it's targeting the privileged SCC
        if let Some(role_ref) = crb_obj.data.get("roleRef") {
            assert_eq!(
                role_ref.get("name").and_then(|v| v.as_str()),
                Some("system:openshift:scc:privileged")
            );
        }

        // Second resource should be the Pod
        let pod_obj = &bundle.resources[1];
        assert_eq!(pod_obj.metadata.name.as_deref(), Some("test-bundle-ocp"));
        assert_eq!(pod_obj.metadata.namespace.as_deref(), Some("test-ns"));
    }

    #[test]
    fn test_build_privileged_bundle_k3s_distribution() {
        let bundle = build_privileged_bundle(
            PrivilegedPodConfig {
                name: "test-bundle-k3s".to_string(),
                namespace: "test-ns".to_string(),
                node_name: "node-1".to_string(),
                container_name: "test-container".to_string(),
                command: vec!["echo".to_string(), "hello".to_string()],
                ..Default::default()
            },
            &KubernetesDistribution::K3sFamily,
        );

        // For K3s, only the Pod should be in the bundle (no special SCC)
        assert_eq!(bundle.resources.len(), 1);

        let pod_obj = &bundle.resources[0];
        assert_eq!(pod_obj.metadata.name.as_deref(), Some("test-bundle-k3s"));
    }

    #[test]
    fn test_pod_yaml_rendering_expected() {
        let pod = build_privileged_pod(
            PrivilegedPodConfig {
                name: "pod_name".to_string(),
                namespace: "pod_namespace".to_string(),
                node_name: "node name".to_string(),
                container_name: "container name".to_string(),
                command: vec!["command".to_string(), "argument".to_string()],
                host_pid: true,
                host_network: true,
                ..Default::default()
            },
            &KubernetesDistribution::Default,
        );

        assert_eq!(
            &serde_yaml::to_string(&pod).unwrap(),
            "apiVersion: v1
kind: Pod
metadata:
  name: pod_name
  namespace: pod_namespace
spec:
  containers:
  - command:
    - command
    - argument
    image: hub.nationtech.io/redhat/ubi10:latest
    name: container name
    securityContext:
      privileged: true
    volumeMounts: []
  hostNetwork: true
  hostPID: true
  nodeName: node name
  restartPolicy: Never
  volumes: []
"
        );
    }

    #[test]
    fn test_pod_yaml_rendering_openshift() {
        let pod = build_privileged_pod(
            PrivilegedPodConfig {
                name: "pod_name".to_string(),
                namespace: "pod_namespace".to_string(),
                node_name: "node name".to_string(),
                container_name: "container name".to_string(),
                command: vec!["command".to_string(), "argument".to_string()],
                host_pid: true,
                host_network: true,
                ..Default::default()
            },
            &KubernetesDistribution::OpenshiftFamily,
        );

        assert_eq!(
            &serde_yaml::to_string(&pod).unwrap(),
            "apiVersion: v1
kind: Pod
metadata:
  annotations:
    openshift.io/required-scc: privileged
    openshift.io/scc: privileged
  name: pod_name
  namespace: pod_namespace
spec:
  containers:
  - command:
    - command
    - argument
    image: hub.nationtech.io/redhat/ubi10:latest
    name: container name
    securityContext:
      privileged: true
    volumeMounts: []
  hostNetwork: true
  hostPID: true
  nodeName: node name
  restartPolicy: Never
  volumes: []
"
        );
    }
}