vault-hier/src/vault_init.rs

use anyhow::{Context, Result};
use reqwest::Client;
use std::{
    env,
    fs,
};
use tracing::{info, warn, error, debug};

use crate::vault_setup::VaultClient;

/// Initialize and unseal the Vault, returning the root token for further operations
pub async fn initialize_vault(vault_addr: &str) -> Result<String> {
    let client = Client::new();

    // Wait for Vault to be available
    VaultClient::wait_for_vault(vault_addr).await?;

    // Display Vault health status
    let health_url = format!("{}/v1/sys/health?standbyok=true&sealedok=true&uninitok=true", vault_addr);
    match client.get(&health_url).send().await {
        Ok(response) => {
            if response.status().is_success() {
                let status_text = response.text().await?;
                info!("Vault status: {}", status_text);
            }
        },
        Err(e) => warn!("Error getting Vault status: {}", e),
    }

    // First check if Vault is already initialized
    let initialized = VaultClient::check_init_status(&client, vault_addr).await?;
    let mut root_token = String::new();

    if initialized {
        info!("Vault is already initialized.");

        // Check if Vault is sealed
        let seal_status = VaultClient::check_seal_status(&client, vault_addr).await?;

        if seal_status.sealed {
            info!("Vault is sealed. Looking for unseal keys...");

            // Try to load unseal keys from environment variables
            let mut unseal_keys = Vec::new();
            for i in 1..=5 {
                match env::var(format!("VAULT_UNSEAL_KEY_{}", i)) {
                    Ok(key) => {
                        info!("Found unseal key {} from environment", i);
                        unseal_keys.push(key);
                    },
                    Err(_) => {
                        debug!("Unseal key {} not found in environment", i);
                    }
                }
            }

            // If we have unseal keys, try to unseal
            if !unseal_keys.is_empty() {
                info!("Found {} unseal keys. Attempting to unseal...", unseal_keys.len());
                VaultClient::unseal_vault(&client, vault_addr, &unseal_keys).await?;
            } else {
                warn!("No unseal keys found. Vault remains sealed.");
                info!("To unseal, set VAULT_UNSEAL_KEY_1, VAULT_UNSEAL_KEY_2, etc. environment variables.");
            }
        } else {
            info!("Vault is already unsealed.");
        }

        // Try to load root token from environment or credentials file
        match env::var("VAULT_TOKEN") {
            Ok(token) => {
                info!("Found root token from environment");
                root_token = token;
            },
            Err(_) => {
                // Try to load from credentials file
                if let Ok(contents) = fs::read_to_string("vault-credentials.json") {
                    if let Ok(creds) = serde_json::from_str::<serde_json::Value>(&contents) {
                        if let Some(token) = creds["root_token"].as_str() {
                            info!("Found root token from credentials file");
                            root_token = token.to_string();
                        }
                    }
                }
            }
        }

        if root_token.is_empty() {
            error!("Unable to find root token. Please set VAULT_TOKEN environment variable or provide vault-credentials.json file.");
            anyhow::bail!("Unable to find root token. Please set VAULT_TOKEN environment variable or provide vault-credentials.json file.");
        }
    } else {
        // Initialize Vault
        info!("Vault is not initialized. Proceeding with initialization...");
        let init_response = VaultClient::init_vault(&client, vault_addr).await?;

        // Save credentials to files
        info!("Saving credentials to files...");
        let current_dir = std::env::current_dir().context("Failed to get current directory")?;

        // Save as JSON (new format)
        let json_path = current_dir.join("vault-credentials.json");
        VaultClient::save_credentials(&init_response, json_path.to_str().unwrap())?;
        info!("JSON credentials saved to: {}", json_path.display());

        // Save as text (for backward compatibility)
        let text_path = current_dir.join("vault-credentials.txt");
        VaultClient::save_credentials(&init_response, text_path.to_str().unwrap())?;
        info!("Text credentials saved to: {}", text_path.display());

        // Also save to /app/data as a backup for Docker volume mounting
        if let Ok(()) = std::fs::create_dir_all("/app/data") {
            let docker_json_path = "/app/data/vault-credentials.json";
            VaultClient::save_credentials(&init_response, docker_json_path)?;
            info!("Backup JSON credentials saved to Docker volume at: {}", docker_json_path);

            let docker_text_path = "/app/data/vault-credentials.txt";
            VaultClient::save_credentials(&init_response, docker_text_path)?;
            info!("Backup text credentials saved to Docker volume at: {}", docker_text_path);
        }

        info!("=========================================");
        info!("IMPORTANT: SAVE THESE CREDENTIALS SECURELY");
        info!("=========================================");
        info!("Root Token: {}", init_response.root_token);
        info!("Unseal Keys (first 3 of 5 needed to unseal):");
        for (i, key) in init_response.keys_base64.iter().enumerate() {
            info!("Key {}: {}", i + 1, key);
        }
        info!("=========================================");

        // Unseal Vault using the first three keys
        let unseal_keys = init_response.keys_base64.iter()
            .take(3)  // We only need threshold number of keys (3)
            .cloned()
            .collect::<Vec<String>>();

        VaultClient::unseal_vault(&client, vault_addr, &unseal_keys).await?;

        info!("Vault initialization and unseal complete!");

        // Set root token
        root_token = init_response.root_token;
    }

    // Look for any existing credentials and copy them to the mounted volume
    if let Ok(metadata) = std::fs::metadata("vault-credentials.json") {
        if metadata.is_file() {
            info!("Found JSON credentials file, ensuring it's saved to Docker volume...");
            if let Ok(()) = std::fs::create_dir_all("/app/data") {
                match std::fs::copy("vault-credentials.json", "/app/data/vault-credentials.json") {
                    Ok(_) => info!("JSON credentials saved to Docker volume"),
                    Err(e) => warn!("Failed to copy JSON credentials: {}", e),
                }
            }
        }
    }

    if let Ok(metadata) = std::fs::metadata("vault-credentials.txt") {
        if metadata.is_file() {
            info!("Found text credentials file, ensuring it's saved to Docker volume...");
            if let Ok(()) = std::fs::create_dir_all("/app/data") {
                match std::fs::copy("vault-credentials.txt", "/app/data/vault-credentials.txt") {
                    Ok(_) => info!("Text credentials saved to Docker volume"),
                    Err(e) => warn!("Failed to copy text credentials: {}", e),
                }
            }
        }
    }

    info!("Vault setup complete!");
    Ok(root_token)
}