lmrc_k3s/

lib.rs

//! # K3s Manager
//!
//! A comprehensive Rust library for managing K3s Kubernetes clusters via SSH.

pub mod adapter;
mod error;
mod state;

pub use adapter::K3sAdapter;
pub use error::{K3sError, Result};
pub use state::{
    ClusterState, ClusterStatus, DesiredClusterConfig, NodeInfo, NodeStatus, ReconciliationPlan,
    WorkerConfig, WorkerNode,
};

use lmrc_ssh::{AuthMethod, SshClient};
use std::path::Path;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tracing::{debug, info};

/// Main K3s cluster manager
pub struct K3sManager {
    k3s_version: String,
    token: String,
    disable_components: Vec<String>,
    ssh_username: String,
    ssh_key_path: Option<String>,
    ssh_password: Option<String>,
    state: Arc<RwLock<Option<ClusterState>>>,
}

impl K3sManager {
    /// Create a new K3sManager
    pub fn new(k3s_version: String, token: String, disable_components: Vec<String>) -> Self {
        Self {
            k3s_version,
            token,
            disable_components,
            ssh_username: "root".to_string(),
            ssh_key_path: None,
            ssh_password: None,
            state: Arc::new(RwLock::new(None)),
        }
    }

    /// Create a new builder
    pub fn builder() -> K3sManagerBuilder {
        K3sManagerBuilder::default()
    }

    /// Execute SSH command
    pub async fn execute_ssh(&self, host: &str, command: &str) -> Result<String> {
        debug!("Executing on {}: {}", host, command);

        // Determine authentication method
        let auth = if let Some(ref password) = self.ssh_password {
            AuthMethod::Password {
                username: self.ssh_username.clone(),
                password: password.clone(),
            }
        } else {
            // Use public key authentication (default to project-local .ssh/id_rsa)
            let key_path = self.ssh_key_path.clone().unwrap_or_else(|| {
                std::env::var("SSH_KEY_PATH")
                    .unwrap_or_else(|_| ".ssh/id_rsa".to_string())
            });

            AuthMethod::PublicKey {
                username: self.ssh_username.clone(),
                private_key_path: key_path,
                passphrase: None,
            }
        };

        // Create SSH client and connect
        let mut client = SshClient::new(host, 22)
            .map_err(|e| K3sError::ConnectionError {
                host: host.to_string(),
                message: format!("Failed to create SSH client: {}", e),
            })?
            .with_auth(auth)
            .connect()
            .map_err(|e| K3sError::ConnectionError {
                host: host.to_string(),
                message: format!("Failed to connect: {}", e),
            })?;

        // Execute command
        let output = client
            .execute(command)
            .map_err(|e| K3sError::CommandError {
                host: host.to_string(),
                command: command.to_string(),
                exit_code: None,
                message: format!("Execution failed: {}", e),
            })?;

        // Check exit status
        if output.exit_status != 0 {
            return Err(K3sError::CommandError {
                host: host.to_string(),
                command: command.to_string(),
                exit_code: Some(output.exit_status),
                message: format!("stderr: {}", output.stderr),
            });
        }

        Ok(output.stdout)
    }

    /// Check if K3s is installed
    pub async fn is_installed(&self, server_ip: &str) -> Result<bool> {
        match self.execute_ssh(server_ip, "which k3s").await {
            Ok(output) => Ok(!output.trim().is_empty()),
            Err(_) => Ok(false),
        }
    }

    /// Install K3s on master
    pub async fn install_master(
        &self,
        server_ip: &str,
        private_ip: Option<&str>,
        force: bool,
    ) -> Result<()> {
        info!("Installing K3s on master node {}", server_ip);

        if !force && self.is_installed(server_ip).await? {
            return Err(K3sError::AlreadyInstalledError {
                node: server_ip.to_string(),
            });
        }

        {
            let mut state = self.state.write().unwrap();
            *state = Some(ClusterState::new(
                server_ip.to_string(),
                private_ip.map(String::from),
                self.k3s_version.clone(),
                self.token.clone(),
                self.disable_components.clone(),
            ));
        }

        let mut install_cmd = if self.k3s_version == "latest" || self.k3s_version.is_empty() {
            format!(
                "curl -sfL https://get.k3s.io | K3S_TOKEN={} sh -s - server --cluster-init",
                self.token
            )
        } else {
            format!(
                "curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION={} K3S_TOKEN={} sh -s - server --cluster-init",
                self.k3s_version, self.token
            )
        };

        if !self.disable_components.is_empty() {
            let disable_args = self
                .disable_components
                .iter()
                .map(|c| format!("--disable {}", c))
                .collect::<Vec<_>>()
                .join(" ");
            install_cmd.push_str(&format!(" {}", disable_args));
        }

        if let Some(priv_ip) = private_ip {
            install_cmd.push_str(&format!(" --node-ip={} --tls-san={}", priv_ip, priv_ip));
        }

        self.execute_ssh(server_ip, &install_cmd).await?;

        // Note: Firewall is now configured via Hetzner's native firewall API
        // during server provisioning, not via UFW. See libs/lmrc-pipeline/src/steps/provision.rs

        self.wait_for_k3s_ready(server_ip, 30, Duration::from_secs(10))
            .await?;

        info!("K3s master installation complete");
        Ok(())
    }

    /// Join worker to cluster
    pub async fn join_worker(
        &self,
        worker_ip: &str,
        master_ip: &str,
        worker_private_ip: Option<&str>,
    ) -> Result<()> {
        info!("Joining worker {} to cluster", worker_ip);

        // Determine which IP to use for cluster communication
        // Use master's private IP from cluster state if available, otherwise fall back to public IP
        let master_cluster_ip = {
            let state = self.state.read().unwrap();
            state.as_ref()
                .and_then(|s| s.master_private_ip.clone())
                .unwrap_or_else(|| master_ip.to_string())
        };

        info!("Using master IP {} for cluster communication", master_cluster_ip);

        // Pre-flight checks: Test connectivity from worker to master
        debug!("Running pre-flight connectivity checks...");

        // Check if worker can reach master on port 6443
        debug!("Testing K3s API connectivity from worker to master...");
        let connectivity_test = format!(
            "timeout 10 nc -zv {} 6443 2>&1 || echo 'CONNECTION_FAILED'",
            master_cluster_ip
        );
        let connectivity_result = self.execute_ssh(worker_ip, &connectivity_test).await?;

        if connectivity_result.contains("CONNECTION_FAILED") || connectivity_result.contains("timed out") {
            return Err(K3sError::ConnectionError {
                host: worker_ip.to_string(),
                message: format!(
                    "Worker cannot reach master on port 6443. This usually indicates a firewall issue. \
                     Connectivity test output: {}",
                    connectivity_result
                ),
            });
        }
        debug!("K3s API connectivity check passed");

        // Check if master K3s service is actually running
        debug!("Verifying master K3s service is running...");
        let service_check = "systemctl is-active k3s 2>&1";
        let service_result = self.execute_ssh(master_ip, service_check).await?;

        if !service_result.trim().contains("active") {
            return Err(K3sError::ClusterNotReadyError {
                attempts: 1,
                message: format!(
                    "Master K3s service is not running. Service status: {}",
                    service_result
                ),
            });
        }
        debug!("Master K3s service is running");

        {
            let mut state = self.state.write().unwrap();
            if let Some(ref mut cluster_state) = *state {
                cluster_state
                    .add_worker(worker_ip.to_string(), worker_private_ip.map(String::from));
            }
        }

        let mut install_cmd = if self.k3s_version == "latest" || self.k3s_version.is_empty() {
            format!(
                "curl -sfL https://get.k3s.io | K3S_URL=https://{}:6443 K3S_TOKEN={} sh -s -",
                master_cluster_ip, self.token
            )
        } else {
            format!(
                "curl -sfL https://get.k3s.io | INSTALL_K3S_VERSION={} K3S_URL=https://{}:6443 K3S_TOKEN={} sh -s -",
                self.k3s_version, master_cluster_ip, self.token
            )
        };

        if let Some(priv_ip) = worker_private_ip {
            install_cmd.push_str(&format!(" --node-ip={}", priv_ip));
        }

        info!("Running K3s agent installation on worker (this may take 2-5 minutes)...");
        self.execute_ssh(worker_ip, &install_cmd).await?;

        info!("K3s agent installation completed, waiting for node to be ready...");
        tokio::time::sleep(Duration::from_secs(5)).await;

        {
            let mut state = self.state.write().unwrap();
            if let Some(ref mut cluster_state) = *state {
                cluster_state.update_worker_status(worker_ip, NodeStatus::Ready);
                cluster_state.update_cluster_status();
            }
        }

        info!("Worker joined successfully");
        Ok(())
    }

    /// Download kubeconfig
    pub async fn download_kubeconfig<P: AsRef<Path>>(
        &self,
        master_ip: &str,
        output_path: P,
    ) -> Result<()> {
        info!("Downloading kubeconfig");

        let content = self
            .execute_ssh(master_ip, "cat /etc/rancher/k3s/k3s.yaml")
            .await?;
        let kubeconfig = content.replace("127.0.0.1", master_ip);

        std::fs::write(output_path.as_ref(), kubeconfig).map_err(|e| K3sError::IoError {
            operation: "write kubeconfig".to_string(),
            source: e,
        })?;

        Ok(())
    }

    /// Get cluster nodes
    pub async fn get_nodes(&self, master_ip: &str) -> Result<String> {
        self.execute_ssh(master_ip, "kubectl get nodes").await
    }

    /// Wait for K3s ready
    async fn wait_for_k3s_ready(
        &self,
        server_ip: &str,
        max_attempts: u32,
        wait_duration: Duration,
    ) -> Result<()> {
        for attempt in 1..=max_attempts {
            match self.execute_ssh(server_ip, "kubectl get nodes").await {
                Ok(_) => return Ok(()),
                Err(e) => {
                    if attempt == max_attempts {
                        return Err(K3sError::ClusterNotReadyError {
                            attempts: max_attempts,
                            message: format!("{}", e),
                        });
                    }
                    tokio::time::sleep(wait_duration).await;
                }
            }
        }
        Ok(())
    }

    /// Uninstall K3s
    pub async fn uninstall(&self, server_ip: &str, is_master: bool) -> Result<()> {
        let script = if is_master {
            "/usr/local/bin/k3s-uninstall.sh"
        } else {
            "/usr/local/bin/k3s-agent-uninstall.sh"
        };

        let _ = self.execute_ssh(server_ip, script).await;
        Ok(())
    }

    /// Get cluster state
    pub fn get_cluster_state(&self) -> Option<ClusterState> {
        self.state.read().unwrap().clone()
    }

    /// Refresh cluster state
    pub async fn refresh_cluster_state(&self, master_ip: &str) -> Result<()> {
        let _ = self.get_nodes(master_ip).await?;

        let mut state = self.state.write().unwrap();
        if let Some(ref mut cluster_state) = *state {
            cluster_state.update_cluster_status();
        }
        Ok(())
    }

    /// Parse kubectl get nodes output into NodeInfo list
    pub async fn get_node_info_list(&self, master_ip: &str) -> Result<Vec<NodeInfo>> {
        let output = self.get_nodes(master_ip).await?;

        let nodes: Vec<NodeInfo> = output
            .lines()
            .skip(1) // Skip header line
            .filter_map(NodeInfo::parse_from_kubectl_line)
            .collect();

        Ok(nodes)
    }

    /// Create a reconciliation plan by comparing desired state with current state
    ///
    /// This method checks what needs to be done to reach the desired cluster configuration.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use lmrc_k3s::{K3sManager, DesiredClusterConfig, WorkerConfig};
    /// # async fn example() -> Result<(), lmrc_k3s::K3sError> {
    /// # let manager = K3sManager::builder().token("token").build();
    /// let desired = DesiredClusterConfig {
    ///     master_ip: "192.168.1.10".to_string(),
    ///     master_private_ip: None,
    ///     workers: vec![
    ///         WorkerConfig::new("192.168.1.11".to_string()),
    ///         WorkerConfig::new("192.168.1.12".to_string()),
    ///     ],
    ///     k3s_version: "v1.28.5+k3s1".to_string(),
    ///     token: "my-token".to_string(),
    ///     disabled_components: vec![],
    /// };
    ///
    /// let plan = manager.plan_reconciliation(&desired).await?;
    /// println!("Plan: {}", plan.summary());
    /// # Ok(())
    /// # }
    /// ```
    pub async fn plan_reconciliation(
        &self,
        desired: &DesiredClusterConfig,
    ) -> Result<ReconciliationPlan> {
        // Check if master is installed
        let master_installed = self.is_installed(&desired.master_ip).await?;

        // Get current worker state
        let current_state = self.get_cluster_state();
        let current_workers: Vec<String> = current_state
            .as_ref()
            .map(|s| s.workers.iter().map(|w| w.ip.clone()).collect())
            .unwrap_or_default();

        let desired_workers: Vec<String> = desired.workers.iter().map(|w| w.ip.clone()).collect();

        // Workers to add = in desired but not in current
        let workers_to_add: Vec<WorkerConfig> = desired
            .workers
            .iter()
            .filter(|w| !current_workers.contains(&w.ip))
            .cloned()
            .collect();

        // Workers to remove = in current but not in desired
        let workers_to_remove: Vec<String> = current_workers
            .into_iter()
            .filter(|ip| !desired_workers.contains(ip))
            .collect();

        Ok(ReconciliationPlan {
            workers_to_add,
            workers_to_remove,
            install_master: !master_installed,
        })
    }

    /// Apply a reconciliation plan to reach the desired state
    ///
    /// This method executes the plan to add/remove nodes as needed.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use lmrc_k3s::{K3sManager, DesiredClusterConfig, WorkerConfig};
    /// # async fn example() -> Result<(), lmrc_k3s::K3sError> {
    /// # let manager = K3sManager::builder().token("token").build();
    /// # let desired = DesiredClusterConfig {
    /// #     master_ip: "192.168.1.10".to_string(),
    /// #     master_private_ip: None,
    /// #     workers: vec![],
    /// #     k3s_version: "v1.28.5+k3s1".to_string(),
    /// #     token: "my-token".to_string(),
    /// #     disabled_components: vec![],
    /// # };
    /// let plan = manager.plan_reconciliation(&desired).await?;
    ///
    /// if plan.has_changes() {
    ///     println!("Applying changes: {}", plan.summary());
    ///     manager.apply_reconciliation(&desired, &plan).await?;
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn apply_reconciliation(
        &self,
        desired: &DesiredClusterConfig,
        plan: &ReconciliationPlan,
    ) -> Result<()> {
        // Install master if needed
        if plan.install_master {
            info!("Installing master as part of reconciliation");
            self.install_master(
                &desired.master_ip,
                desired.master_private_ip.as_deref(),
                false,
            )
            .await?;
        }

        // Remove workers first
        for worker_ip in &plan.workers_to_remove {
            info!("Removing worker node: {}", worker_ip);
            self.remove_worker(&desired.master_ip, worker_ip).await?;
        }

        // Add new workers
        for worker in &plan.workers_to_add {
            info!("Adding worker node: {}", worker.ip);
            self.join_worker(&worker.ip, &desired.master_ip, worker.private_ip.as_deref())
                .await?;
        }

        Ok(())
    }

    /// Reconcile cluster to match desired configuration
    ///
    /// This is a convenience method that combines `plan_reconciliation` and `apply_reconciliation`.
    /// It will automatically add or remove workers to match the desired state.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use lmrc_k3s::{K3sManager, DesiredClusterConfig, WorkerConfig};
    /// # async fn example() -> Result<(), lmrc_k3s::K3sError> {
    /// let manager = K3sManager::builder()
    ///     .token("my-token")
    ///     .build();
    ///
    /// // Define desired state
    /// let desired = DesiredClusterConfig {
    ///     master_ip: "192.168.1.10".to_string(),
    ///     master_private_ip: None,
    ///     workers: vec![
    ///         WorkerConfig::new("192.168.1.11".to_string()),
    ///         WorkerConfig::new("192.168.1.12".to_string()),
    ///         WorkerConfig::new("192.168.1.13".to_string()), // New worker
    ///     ],
    ///     k3s_version: "v1.28.5+k3s1".to_string(),
    ///     token: "my-token".to_string(),
    ///     disabled_components: vec!["traefik".to_string()],
    /// };
    ///
    /// // Reconcile - will add worker .13 if not present
    /// manager.reconcile(&desired).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn reconcile(&self, desired: &DesiredClusterConfig) -> Result<()> {
        let plan = self.plan_reconciliation(desired).await?;

        if !plan.has_changes() {
            info!("Cluster is already in desired state");
            return Ok(());
        }

        info!("Reconciliation plan: {}", plan.summary());
        self.apply_reconciliation(desired, &plan).await?;

        info!("Reconciliation complete");
        Ok(())
    }

    /// Remove a worker node from the cluster
    ///
    /// This method drains the node, deletes it from the cluster, and uninstalls K3s.
    ///
    /// # Arguments
    ///
    /// * `master_ip` - The IP of the master node
    /// * `worker_ip` - The IP of the worker node to remove
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use lmrc_k3s::K3sManager;
    /// # async fn example() -> Result<(), lmrc_k3s::K3sError> {
    /// # let manager = K3sManager::builder().token("token").build();
    /// manager.remove_worker("192.168.1.10", "192.168.1.11").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn remove_worker(&self, master_ip: &str, worker_ip: &str) -> Result<()> {
        info!("Removing worker node: {}", worker_ip);

        // Get the node name from kubectl
        let nodes = self.get_node_info_list(master_ip).await?;
        let node_name = nodes
            .iter()
            .find(|n| {
                n.internal_ip.as_ref() == Some(&worker_ip.to_string())
                    || n.external_ip.as_ref() == Some(&worker_ip.to_string())
                    || n.name.contains(worker_ip)
            })
            .map(|n| n.name.clone());

        if let Some(name) = node_name {
            // Drain the node
            info!("Draining node: {}", name);
            let drain_cmd = format!(
                "kubectl drain {} --ignore-daemonsets --delete-emptydir-data --force --timeout=60s",
                name
            );
            let _ = self.execute_ssh(master_ip, &drain_cmd).await;

            // Delete from cluster
            info!("Deleting node from cluster: {}", name);
            let delete_cmd = format!("kubectl delete node {}", name);
            let _ = self.execute_ssh(master_ip, &delete_cmd).await;
        }

        // Uninstall K3s from the worker
        info!("Uninstalling K3s from worker: {}", worker_ip);
        self.uninstall(worker_ip, false).await?;

        // Update state
        {
            let mut state = self.state.write().unwrap();
            if let Some(ref mut cluster_state) = *state {
                cluster_state.workers.retain(|w| w.ip != worker_ip);
                cluster_state.update_cluster_status();
            }
        }

        info!("Worker node removed successfully");
        Ok(())
    }
}

/// Builder for K3sManager
#[derive(Default)]
pub struct K3sManagerBuilder {
    k3s_version: Option<String>,
    token: Option<String>,
    disable_components: Vec<String>,
    ssh_username: Option<String>,
    ssh_key_path: Option<String>,
    ssh_password: Option<String>,
}

impl K3sManagerBuilder {
    /// Set K3s version
    pub fn version(mut self, version: impl Into<String>) -> Self {
        self.k3s_version = Some(version.into());
        self
    }

    /// Set cluster token (required)
    pub fn token(mut self, token: impl Into<String>) -> Self {
        self.token = Some(token.into());
        self
    }

    /// Set components to disable
    pub fn disable(mut self, components: Vec<String>) -> Self {
        self.disable_components = components;
        self
    }

    /// Set SSH username (default: "root")
    pub fn ssh_username(mut self, username: impl Into<String>) -> Self {
        self.ssh_username = Some(username.into());
        self
    }

    /// Set SSH private key path (default: ~/.ssh/id_rsa)
    pub fn ssh_key_path(mut self, path: impl Into<String>) -> Self {
        self.ssh_key_path = Some(path.into());
        self
    }

    /// Set SSH password for password-based authentication
    /// If not set, public key authentication will be used
    pub fn ssh_password(mut self, password: impl Into<String>) -> Self {
        self.ssh_password = Some(password.into());
        self
    }

    /// Build the K3sManager instance
    pub fn build(self) -> K3sManager {
        K3sManager {
            k3s_version: self
                .k3s_version
                .unwrap_or_else(|| "v1.28.5+k3s1".to_string()),
            token: self.token.expect("Token must be set"),
            disable_components: self.disable_components,
            ssh_username: self.ssh_username.unwrap_or_else(|| "root".to_string()),
            ssh_key_path: self.ssh_key_path,
            ssh_password: self.ssh_password,
            state: Arc::new(RwLock::new(None)),
        }
    }
}