foxhunt/vault-migration/migration-scripts/populate-vault.rs

//! Migration script to populate HashiCorp Vault with existing secrets from environment variables
//!
//! This script reads secrets from the current environment and migrates them to Vault
//! following the organized structure defined in vault-structure.md

use anyhow::{Context, Result};
use clap::{Arg, Command};
use serde_json::{json, Map, Value};
use std::collections::HashMap;
use std::env;
use tracing::{error, info, warn};
use url::Url;

// Re-export vault client types
use foxhunt_vault_client::{
    auth::{AppRoleTokenProvider, KubernetesTokenProvider, StaticTokenProvider},
    FoxhuntVaultClient, VaultClientConfig, WriteOptions,
};

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .init();

    let matches = Command::new("vault-migration")
        .about("Migrates Foxhunt secrets from environment variables to HashiCorp Vault")
        .arg(
            Arg::new("vault-url")
                .long("vault-url")
                .value_name("URL")
                .help("Vault server URL")
                .default_value("https://vault.foxhunt.local:8200"),
        )
        .arg(
            Arg::new("environment")
                .long("environment")
                .short('e')
                .value_name("ENV")
                .help("Target environment (development, staging, production)")
                .default_value("development"),
        )
        .arg(
            Arg::new("service")
                .long("service")
                .short('s')
                .value_name("SERVICE")
                .help("Service name for secret namespacing")
                .default_value("trading-service"),
        )
        .arg(
            Arg::new("auth-method")
                .long("auth-method")
                .value_name("METHOD")
                .help("Authentication method (approle, kubernetes, token)")
                .default_value("token"),
        )
        .arg(
            Arg::new("role-id")
                .long("role-id")
                .value_name("ROLE_ID")
                .help("AppRole role ID (required for AppRole auth)")
                .required_if_eq("auth-method", "approle"),
        )
        .arg(
            Arg::new("secret-id")
                .long("secret-id")
                .value_name("SECRET_ID")
                .help("AppRole secret ID (required for AppRole auth)")
                .required_if_eq("auth-method", "approle"),
        )
        .arg(
            Arg::new("k8s-role")
                .long("k8s-role")
                .value_name("ROLE")
                .help("Kubernetes role name (required for Kubernetes auth)")
                .required_if_eq("auth-method", "kubernetes"),
        )
        .arg(
            Arg::new("vault-token")
                .long("vault-token")
                .value_name("TOKEN")
                .help("Static Vault token (required for token auth)")
                .required_if_eq("auth-method", "token"),
        )
        .arg(
            Arg::new("dry-run")
                .long("dry-run")
                .help("Show what would be migrated without actually writing to Vault")
                .action(clap::ArgAction::SetTrue),
        )
        .get_matches();

    let vault_url = matches.get_one::<String>("vault-url").unwrap();
    let environment = matches.get_one::<String>("environment").unwrap();
    let service = matches.get_one::<String>("service").unwrap();
    let auth_method = matches.get_one::<String>("auth-method").unwrap();
    let dry_run = matches.get_flag("dry-run");

    info!("Starting Foxhunt secret migration to Vault");
    info!("Vault URL: {}", vault_url);
    info!("Environment: {}", environment);
    info!("Service: {}", service);
    info!("Auth method: {}", auth_method);
    info!("Dry run: {}", dry_run);

    // Create Vault client configuration
    let config = VaultClientConfig {
        vault_url: vault_url.clone(),
        environment: environment.clone(),
        service_name: service.clone(),
        ..Default::default()
    };

    // Create token provider based on auth method
    let token_provider: Box<dyn foxhunt_vault_client::TokenProvider + Send + Sync> = match auth_method.as_str() {
        "approle" => {
            let role_id = matches.get_one::<String>("role-id").unwrap().clone();
            let secret_id = matches.get_one::<String>("secret-id").unwrap().clone();
            let vault_url = Url::parse(vault_url)?;
            let http_client = reqwest::Client::new();
            Box::new(AppRoleTokenProvider::new(vault_url, http_client, role_id, secret_id))
        }
        "kubernetes" => {
            let k8s_role = matches.get_one::<String>("k8s-role").unwrap().clone();
            let vault_url = Url::parse(vault_url)?;
            let http_client = reqwest::Client::new();
            Box::new(KubernetesTokenProvider::new(vault_url, http_client, k8s_role, None))
        }
        "token" => {
            let vault_token = matches.get_one::<String>("vault-token").unwrap().clone();
            Box::new(StaticTokenProvider::new(vault_token))
        }
        _ => {
            error!("Invalid auth method: {}", auth_method);
            std::process::exit(1);
        }
    };

    // Create Vault client
    let vault_client = FoxhuntVaultClient::new(config, token_provider)
        .await
        .context("Failed to create Vault client")?;

    // Perform health check
    vault_client.health_check().await.context("Vault health check failed")?;
    info!("Vault client initialized and health check passed");

    // Collect secrets from environment
    let secrets = collect_secrets_from_environment();
    info!("Collected {} secret categories from environment", secrets.len());

    if dry_run {
        info!("DRY RUN - Showing secrets that would be migrated:");
        for (category, secret_map) in &secrets {
            info!("Category: {}", category);
            for (path, data) in secret_map {
                info!("  Path: {} (fields: {})", path, data.keys().count());
                for key in data.keys() {
                    info!("    - {}", key);
                }
            }
        }
        info!("DRY RUN complete. Use --dry-run=false to perform actual migration.");
        return Ok(());
    }

    // Migrate secrets to Vault
    let mut total_migrated = 0;
    let mut failed_migrations = 0;

    for (category, secret_map) in secrets {
        info!("Migrating {} category with {} secrets", category, secret_map.len());

        for (path, data) in secret_map {
            match migrate_secret(&vault_client, &path, data).await {
                Ok(()) => {
                    info!("✅ Successfully migrated: {}", path);
                    total_migrated += 1;
                }
                Err(e) => {
                    error!("❌ Failed to migrate {}: {}", path, e);
                    failed_migrations += 1;
                }
            }
        }
    }

    // Summary
    info!("Migration complete:");
    info!("  Total migrated: {}", total_migrated);
    info!("  Failed migrations: {}", failed_migrations);

    if failed_migrations > 0 {
        error!("Some migrations failed. Please check the logs and retry failed migrations.");
        std::process::exit(1);
    }

    info!("🎉 All secrets successfully migrated to Vault!");
    Ok(())
}

/// Collect secrets from environment variables and organize them by category
fn collect_secrets_from_environment() -> HashMap<String, HashMap<String, HashMap<String, serde_json::Value>>> {
    let mut secrets = HashMap::new();

    // Database secrets
    let mut database_secrets = HashMap::new();
    
    // PostgreSQL
    if let Ok(database_url) = env::var("DATABASE_URL") {
        let mut postgres_data = HashMap::new();
        postgres_data.insert("url".to_string(), json!(database_url));
        
        // Extract components if needed
        if let Ok(parsed) = url::Url::parse(&database_url) {
            if let Some(host) = parsed.host_str() {
                postgres_data.insert("host".to_string(), json!(host));
            }
            if let Some(port) = parsed.port() {
                postgres_data.insert("port".to_string(), json!(port));
            }
            postgres_data.insert("database".to_string(), json!(parsed.path().trim_start_matches('/')));
            if !parsed.username().is_empty() {
                postgres_data.insert("username".to_string(), json!(parsed.username()));
            }
            if let Some(password) = parsed.password() {
                postgres_data.insert("password".to_string(), json!(password));
            }
        }
        
        database_secrets.insert("database/postgresql".to_string(), postgres_data);
    }

    // Redis
    if let Ok(redis_url) = env::var("REDIS_URL") {
        let mut redis_data = HashMap::new();
        redis_data.insert("url".to_string(), json!(redis_url));
        database_secrets.insert("database/redis".to_string(), redis_data);
    }

    // InfluxDB
    if let (Ok(influx_url), Ok(influx_token)) = (env::var("INFLUX_URL"), env::var("INFLUX_TOKEN")) {
        let mut influx_data = HashMap::new();
        influx_data.insert("url".to_string(), json!(influx_url));
        influx_data.insert("token".to_string(), json!(influx_token));
        
        if let Ok(org) = env::var("INFLUX_ORG") {
            influx_data.insert("org".to_string(), json!(org));
        }
        if let Ok(bucket) = env::var("INFLUX_BUCKET") {
            influx_data.insert("bucket".to_string(), json!(bucket));
        }
        
        database_secrets.insert("database/influxdb".to_string(), influx_data);
    }

    // ClickHouse
    if let Ok(clickhouse_url) = env::var("CLICKHOUSE_URL") {
        let mut clickhouse_data = HashMap::new();
        clickhouse_data.insert("url".to_string(), json!(clickhouse_url));
        
        if let Ok(username) = env::var("CLICKHOUSE_USER") {
            clickhouse_data.insert("username".to_string(), json!(username));
        }
        if let Ok(password) = env::var("CLICKHOUSE_PASSWORD") {
            clickhouse_data.insert("password".to_string(), json!(password));
        }
        if let Ok(database) = env::var("CLICKHOUSE_DB") {
            clickhouse_data.insert("database".to_string(), json!(database));
        }
        
        database_secrets.insert("database/clickhouse".to_string(), clickhouse_data);
    }

    if !database_secrets.is_empty() {
        secrets.insert("database".to_string(), database_secrets);
    }

    // API Keys
    let mut api_key_secrets = HashMap::new();
    
    if let Ok(databento_key) = env::var("DATABENTO_API_KEY") {
        let mut databento_data = HashMap::new();
        databento_data.insert("api_key".to_string(), json!(databento_key));
        databento_data.insert("endpoint".to_string(), json!("wss://gateway.databento.com/v2"));
        databento_data.insert("rate_limit".to_string(), json!(10));
        api_key_secrets.insert("api-keys/databento".to_string(), databento_data);
    }

    if let Ok(benzinga_key) = env::var("BENZINGA_API_KEY") {
        let mut benzinga_data = HashMap::new();
        benzinga_data.insert("api_key".to_string(), json!(benzinga_key));
        benzinga_data.insert("endpoint".to_string(), json!("wss://api.benzinga.com/api/v1/news/stream"));
        benzinga_data.insert("rate_limit".to_string(), json!(5));
        api_key_secrets.insert("api-keys/benzinga".to_string(), benzinga_data);
    }

    if let Ok(alpha_vantage_key) = env::var("ALPHA_VANTAGE_API_KEY") {
        let mut alpha_vantage_data = HashMap::new();
        alpha_vantage_data.insert("api_key".to_string(), json!(alpha_vantage_key));
        alpha_vantage_data.insert("endpoint".to_string(), json!("https://www.alphavantage.co"));
        alpha_vantage_data.insert("rate_limit".to_string(), json!(1));
        api_key_secrets.insert("api-keys/alpha-vantage".to_string(), alpha_vantage_data);
    }

    if !api_key_secrets.is_empty() {
        secrets.insert("api-keys".to_string(), api_key_secrets);
    }

    // Authentication secrets
    let mut auth_secrets = HashMap::new();
    
    if let Ok(jwt_secret) = env::var("JWT_SECRET")
        .or_else(|_| env::var("FOXHUNT_JWT_SECRET")) {
        let mut jwt_data = HashMap::new();
        jwt_data.insert("secret".to_string(), json!(jwt_secret));
        jwt_data.insert("issuer".to_string(), json!("foxhunt-hft"));
        jwt_data.insert("audience".to_string(), json!("foxhunt-services"));
        jwt_data.insert("expiration_seconds".to_string(), json!(3600));
        auth_secrets.insert("authentication/jwt".to_string(), jwt_data);
    }

    if let Ok(encryption_key) = env::var("FOXHUNT_ENCRYPTION_KEY") {
        let mut encryption_data = HashMap::new();
        encryption_data.insert("primary_key".to_string(), json!(encryption_key));
        encryption_data.insert("algorithm".to_string(), json!("AES256-GCM"));
        auth_secrets.insert("authentication/encryption".to_string(), encryption_data);
    }

    if !auth_secrets.is_empty() {
        secrets.insert("authentication".to_string(), auth_secrets);
    }

    // Broker credentials
    let mut broker_secrets = HashMap::new();
    
    if let (Ok(ic_username), Ok(ic_password)) = (env::var("ICMARKETS_USERNAME"), env::var("ICMARKETS_PASSWORD")) {
        let mut ic_data = HashMap::new();
        ic_data.insert("username".to_string(), json!(ic_username));
        ic_data.insert("password".to_string(), json!(ic_password));
        ic_data.insert("sender_comp_id".to_string(), json!("FOXHUNT"));
        ic_data.insert("target_comp_id".to_string(), json!("ICMARKETS"));
        ic_data.insert("endpoint".to_string(), json!("fix.icmarkets.com:443"));
        broker_secrets.insert("brokers/icmarkets".to_string(), ic_data);
    }

    if let Ok(ib_host) = env::var("IB_HOST") {
        let mut ib_data = HashMap::new();
        ib_data.insert("host".to_string(), json!(ib_host));
        ib_data.insert("port".to_string(), json!(env::var("IB_PORT").unwrap_or("7497".to_string()).parse::<u16>().unwrap_or(7497)));
        if let Ok(client_id) = env::var("IB_CLIENT_ID") {
            ib_data.insert("client_id".to_string(), json!(client_id.parse::<u32>().unwrap_or(1)));
        }
        if let Ok(account_id) = env::var("IB_ACCOUNT_ID") {
            ib_data.insert("account_id".to_string(), json!(account_id));
        }
        broker_secrets.insert("brokers/interactive-brokers".to_string(), ib_data);
    }

    if !broker_secrets.is_empty() {
        secrets.insert("brokers".to_string(), broker_secrets);
    }

    secrets
}

/// Migrate a single secret to Vault
async fn migrate_secret(
    vault_client: &FoxhuntVaultClient,
    path: &str,
    data: HashMap<String, serde_json::Value>,
) -> Result<()> {
    // Convert the data to the format expected by the Vault client
    let write_options = WriteOptions {
        metadata: Some({
            let mut metadata = HashMap::new();
            metadata.insert("migrated_by".to_string(), "foxhunt-vault-migration".to_string());
            metadata.insert("migrated_at".to_string(), chrono::Utc::now().to_rfc3339());
            metadata.insert("source".to_string(), "environment_variables".to_string());
            metadata
        }),
        ..Default::default()
    };

    vault_client
        .write_secret(path, data, Some(write_options))
        .await
        .context("Failed to write secret to Vault")?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    #[test]
    fn test_collect_secrets_with_database_url() {
        std::env::set_var("DATABASE_URL", "postgresql://user:pass@localhost:5432/foxhunt");
        
        let secrets = collect_secrets_from_environment();
        
        assert!(secrets.contains_key("database"));
        let database_secrets = secrets.get("database").unwrap();
        assert!(database_secrets.contains_key("database/postgresql"));
        
        let postgres_secret = database_secrets.get("database/postgresql").unwrap();
        assert!(postgres_secret.contains_key("url"));
        assert!(postgres_secret.contains_key("host"));
        assert!(postgres_secret.contains_key("port"));
        assert!(postgres_secret.contains_key("database"));
        assert!(postgres_secret.contains_key("username"));
        assert!(postgres_secret.contains_key("password"));
        
        std::env::remove_var("DATABASE_URL");
    }

    #[test]
    fn test_collect_secrets_with_api_keys() {
        std::env::set_var("DATABENTO_API_KEY", "test-databento-key");
        std::env::set_var("BENZINGA_API_KEY", "test-benzinga-key");
        
        let secrets = collect_secrets_from_environment();
        
        assert!(secrets.contains_key("api-keys"));
        let api_secrets = secrets.get("api-keys").unwrap();
        assert!(api_secrets.contains_key("api-keys/databento"));
        assert!(api_secrets.contains_key("api-keys/benzinga"));
        
        std::env::remove_var("DATABENTO_API_KEY");
        std::env::remove_var("BENZINGA_API_KEY");
    }

    #[test]
    fn test_collect_secrets_empty_environment() {
        // Clear relevant environment variables
        let env_vars = ["DATABASE_URL", "REDIS_URL", "DATABENTO_API_KEY", "BENZINGA_API_KEY", "JWT_SECRET"];
        for var in &env_vars {
            std::env::remove_var(var);
        }
        
        let secrets = collect_secrets_from_environment();
        
        // Should return empty or minimal secrets
        assert!(secrets.is_empty() || secrets.values().all(|category| category.is_empty()));
    }
}