harmony/examples/harmony_sso/src/main.rs

use anyhow::Context;
use clap::Parser;
use harmony::inventory::Inventory;
use harmony::modules::k8s::coredns::{CoreDNSRewrite, CoreDNSRewriteScore};
use harmony::modules::openbao::{
    OpenbaoJwtAuth, OpenbaoPolicy, OpenbaoScore, OpenbaoSetupScore, OpenbaoUser,
};
use harmony::modules::zitadel::{
    ZitadelAppType, ZitadelApplication, ZitadelClientConfig, ZitadelScore, ZitadelSetupScore,
};
use harmony::score::Score;
use harmony::topology::{K8sclient, Topology};
use harmony_config::{Config, ConfigClient, EnvSource, StoreSource};
use harmony_k8s::K8sClient;
use harmony_secret::OpenbaoSecretStore;
use k3d_rs::{K3d, PortMapping};
use log::info;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;

const CLUSTER_NAME: &str = "harmony-example";
const ZITADEL_HOST: &str = "sso.harmony.local";
const OPENBAO_HOST: &str = "bao.harmony.local";
const HTTP_PORT: u32 = 8080;
const OPENBAO_NAMESPACE: &str = "openbao";
const OPENBAO_POD: &str = "openbao-0";
const APP_NAME: &str = "harmony-cli";
const PROJECT_NAME: &str = "harmony";

#[derive(Parser)]
#[command(
    name = "harmony-sso",
    about = "Deploy Zitadel + OpenBao on k3d, authenticate via SSO, store config"
)]
struct Args {
    /// Skip Zitadel deployment (OpenBao only, faster iteration)
    #[arg(long)]
    skip_zitadel: bool,

    /// Delete the k3d cluster and exit
    #[arg(long)]
    cleanup: bool,
}

// ---------------------------------------------------------------------------
// Config type stored via SSO-authenticated OpenBao
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq)]
struct SsoExampleConfig {
    team_name: String,
    environment: String,
    max_replicas: u16,
}

impl Default for SsoExampleConfig {
    fn default() -> Self {
        Self {
            team_name: "platform-team".to_string(),
            environment: "staging".to_string(),
            max_replicas: 3,
        }
    }
}

impl Config for SsoExampleConfig {
    const KEY: &'static str = "SsoExampleConfig";
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

fn harmony_data_dir() -> PathBuf {
    directories::BaseDirs::new()
        .map(|dirs| dirs.data_dir().join("harmony"))
        .unwrap_or_else(|| PathBuf::from("/tmp/harmony"))
}

fn create_k3d() -> K3d {
    let base_dir = harmony_data_dir().join("k3d");
    std::fs::create_dir_all(&base_dir).expect("Failed to create k3d data directory");
    K3d::new(base_dir, Some(CLUSTER_NAME.to_string()))
        .with_port_mappings(vec![PortMapping::new(HTTP_PORT, 80)])
}

fn create_topology(k3d: &K3d) -> harmony::topology::K8sAnywhereTopology {
    let context = k3d
        .context_name()
        .unwrap_or_else(|| format!("k3d-{}", CLUSTER_NAME));
    unsafe {
        std::env::set_var("HARMONY_USE_LOCAL_K3D", "false");
        std::env::set_var("HARMONY_AUTOINSTALL", "false");
        std::env::set_var("HARMONY_K8S_CONTEXT", &context);
    }
    harmony::topology::K8sAnywhereTopology::from_env()
}

fn harmony_dev_policy() -> OpenbaoPolicy {
    OpenbaoPolicy {
        name: "harmony-dev".to_string(),
        hcl: r#"path "secret/data/harmony/*" { capabilities = ["create","read","update","delete","list"] }
path "secret/metadata/harmony/*" { capabilities = ["list","read"] }"#
            .to_string(),
    }
}

// ---------------------------------------------------------------------------
// Zitadel deployment (with CNPG retry)
// ---------------------------------------------------------------------------

async fn deploy_zitadel(k3d: &K3d) -> anyhow::Result<()> {
    info!("Deploying Zitadel (this may take several minutes)...");

    let zitadel = ZitadelScore {
        host: ZITADEL_HOST.to_string(),
        zitadel_version: "v4.12.1".to_string(),
        external_secure: false,
        external_port: None,
        ..Default::default()
    };

    let topology = create_topology(k3d);
    topology
        .ensure_ready()
        .await
        .context("Topology init failed")?;

    zitadel
        .interpret(&Inventory::autoload(), &topology)
        .await
        .context("Zitadel deployment failed")?;

    info!("Zitadel deployed successfully");
    Ok(())
}

async fn wait_for_zitadel_ready() -> anyhow::Result<()> {
    info!("Waiting for Zitadel to be ready...");
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(5))
        .build()?;

    for attempt in 1..=90 {
        match client
            .get(format!(
                "http://127.0.0.1:{}/.well-known/openid-configuration",
                HTTP_PORT
            ))
            .header("Host", ZITADEL_HOST)
            .send()
            .await
        {
            Ok(resp) if resp.status().is_success() => {
                info!("Zitadel is ready");
                return Ok(());
            }
            Ok(resp) if attempt % 10 == 0 => {
                info!("Zitadel HTTP {}, attempt {}/90", resp.status(), attempt);
            }
            Err(e) if attempt % 10 == 0 => {
                info!("Zitadel not reachable: {}, attempt {}/90", e, attempt);
            }
            _ => {}
        }
        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
    }

    anyhow::bail!("Timed out waiting for Zitadel")
}

// ---------------------------------------------------------------------------
// Cluster lifecycle
// ---------------------------------------------------------------------------

async fn ensure_k3d_cluster(k3d: &K3d) -> anyhow::Result<()> {
    info!("Ensuring k3d cluster '{}' is running...", CLUSTER_NAME);
    k3d.ensure_installed()
        .await
        .map_err(|e| anyhow::anyhow!("k3d setup failed: {}", e))?;
    info!("k3d cluster '{}' is ready", CLUSTER_NAME);
    Ok(())
}

fn cleanup_cluster(k3d: &K3d) -> anyhow::Result<()> {
    let name = k3d
        .cluster_name()
        .ok_or_else(|| anyhow::anyhow!("No cluster name"))?;
    info!("Deleting k3d cluster '{}'...", name);
    k3d.run_k3d_command(["cluster", "delete", name])
        .map_err(|e| anyhow::anyhow!("{}", e))?;
    info!("Cluster '{}' deleted", name);
    Ok(())
}

async fn cleanup_openbao_webhook(k8s: &K8sClient) -> anyhow::Result<()> {
    use k8s_openapi::api::admissionregistration::v1::MutatingWebhookConfiguration;
    if k8s
        .get_resource::<MutatingWebhookConfiguration>("openbao-agent-injector-cfg", None)
        .await?
        .is_some()
    {
        info!("Deleting conflicting OpenBao webhook...");
        k8s.delete_resource::<MutatingWebhookConfiguration>("openbao-agent-injector-cfg", None)
            .await?;
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
    let args = Args::parse();
    let k3d = create_k3d();

    if args.cleanup {
        return cleanup_cluster(&k3d);
    }

    info!("===========================================");
    info!("Harmony SSO Example");
    info!("===========================================");

    // --- Phase 1: Infrastructure ---

    ensure_k3d_cluster(&k3d).await?;

    let topology = create_topology(&k3d);
    topology
        .ensure_ready()
        .await
        .context("Topology init failed")?;

    let k8s = topology
        .k8s_client()
        .await
        .map_err(|e| anyhow::anyhow!("K8s client: {}", e))?;

    // Deploy + configure OpenBao (no JWT auth yet -- Zitadel isn't up)
    cleanup_openbao_webhook(&k8s).await?;
    OpenbaoScore {
        instance: Default::default(),
        host: OPENBAO_HOST.to_string(),
        openshift: false,
        tls_issuer: None,
    }
    .interpret(&Inventory::autoload(), &topology)
    .await
    .context("OpenBao deploy failed")?;

    OpenbaoSetupScore {
        policies: vec![harmony_dev_policy()],
        users: vec![OpenbaoUser {
            username: "harmony".to_string(),
            password: "harmony-dev-password".to_string(),
            policies: vec!["harmony-dev".to_string()],
        }],
        jwt_auth: None, // Phase 2 adds JWT after Zitadel is ready
        ..Default::default()
    }
    .interpret(&Inventory::autoload(), &topology)
    .await
    .context("OpenBao setup failed")?;

    if args.skip_zitadel {
        info!("=== Skipping Zitadel (--skip-zitadel) ===");
        info!("OpenBao: http://{}:{}", OPENBAO_HOST, HTTP_PORT);
        return Ok(());
    }

    // --- Phase 2: Identity + SSO Wiring ---

    CoreDNSRewriteScore {
        rewrites: vec![
            CoreDNSRewrite {
                hostname: ZITADEL_HOST.to_string(),
                target: "zitadel.zitadel.svc.cluster.local".to_string(),
            },
            CoreDNSRewrite {
                hostname: OPENBAO_HOST.to_string(),
                target: "openbao.openbao.svc.cluster.local".to_string(),
            },
        ],
    }
    .interpret(&Inventory::autoload(), &topology)
    .await
    .context("CoreDNS rewrite failed")?;

    deploy_zitadel(&k3d).await?;
    wait_for_zitadel_ready().await?;

    // Provision Zitadel project + device-code application
    ZitadelSetupScore {
        host: ZITADEL_HOST.to_string(),
        scheme: Default::default(),
        port: None,
        skip_tls: true,
        endpoint: Some(format!("http://127.0.0.1:{HTTP_PORT}")),
        admin_org_id: None,
        namespace: "zitadel".to_string(),
        applications: vec![ZitadelApplication {
            project_name: PROJECT_NAME.to_string(),
            app_name: APP_NAME.to_string(),
            app_type: ZitadelAppType::DeviceCode,
        }],
        api_apps: vec![],
        roles: vec![],
        machine_users: vec![],
    }
    .interpret(&Inventory::autoload(), &topology)
    .await
    .context("Zitadel setup failed")?;

    // Read the client_id from the cache written by ZitadelSetupScore
    let zitadel_config =
        ZitadelClientConfig::load().context("ZitadelSetupScore did not produce a client config")?;
    let client_id = zitadel_config
        .client_id(APP_NAME)
        .context("No client_id for harmony-cli app")?
        .clone();

    info!("Zitadel app '{}' client_id: {}", APP_NAME, client_id);

    // Now configure OpenBao JWT auth with the real client_id
    OpenbaoSetupScore {
        policies: vec![harmony_dev_policy()],
        users: vec![OpenbaoUser {
            username: "harmony".to_string(),
            password: "harmony-dev-password".to_string(),
            policies: vec!["harmony-dev".to_string()],
        }],
        jwt_auth: Some(OpenbaoJwtAuth {
            oidc_discovery_url: format!("http://{}:{}", ZITADEL_HOST, HTTP_PORT),
            bound_issuer: format!("http://{}:{}", ZITADEL_HOST, HTTP_PORT),
            role_name: "harmony-developer".to_string(),
            bound_audiences: client_id.clone(),
            bound_claims_json: String::new(),
            groups_claim: String::new(),
            user_claim: "email".to_string(),
            policies: vec!["harmony-dev".to_string()],
            ttl: "4h".to_string(),
            max_ttl: "24h".to_string(),
        }),
        ..Default::default()
    }
    .interpret(&Inventory::autoload(), &topology)
    .await
    .context("OpenBao JWT auth setup failed")?;

    // --- Phase 3: Config via SSO ---

    info!("===========================================");
    info!("Storing config via SSO-authenticated OpenBao");
    info!("===========================================");

    let _pf = k8s
        .port_forward(OPENBAO_POD, OPENBAO_NAMESPACE, 8200, 8200)
        .await
        .context("Port-forward to OpenBao failed")?;
    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;

    let openbao_url = format!("http://127.0.0.1:{}", _pf.port());
    let sso_url = format!("http://{}:{}", ZITADEL_HOST, HTTP_PORT);

    let store = OpenbaoSecretStore::new(harmony_secret::OpenbaoStoreOptions {
        base_url: openbao_url,
        kv_mount: "secret".to_string(),
        auth_mount: "jwt".to_string(),
        skip_tls: true,
        token: None,
        username: None,
        password: None,
        zitadel_sso_url: Some(sso_url),
        zitadel_client_id: Some(client_id),
        jwt_role: Some("harmony-developer".to_string()),
        jwt_auth_mount: Some("jwt".to_string()),
        zitadel_jwt_bearer: None,
    })
    .await
    .context("SSO authentication failed")?;

    let manager = ConfigClient::new(vec![
        Arc::new(EnvSource) as Arc<dyn harmony_config::ConfigSource>,
        Arc::new(StoreSource::new("harmony".to_string(), store)),
    ]);

    // Try to load existing config (succeeds on re-run)
    match manager.get::<SsoExampleConfig>().await {
        Ok(config) => {
            info!("Config loaded from OpenBao: {:?}", config);
        }
        Err(harmony_config::ConfigError::NotFound { .. }) => {
            info!("No config found, storing default...");
            let config = SsoExampleConfig::default();
            manager.set(&config).await?;
            info!("Config stored: {:?}", config);

            let retrieved: SsoExampleConfig = manager.get().await?;
            info!("Config verified: {:?}", retrieved);
            assert_eq!(config, retrieved);
        }
        Err(e) => return Err(e.into()),
    }

    info!("===========================================");
    info!("Success! Config managed via Zitadel SSO + OpenBao");
    info!("===========================================");
    info!("OpenBao: http://{}:{}", OPENBAO_HOST, HTTP_PORT);
    info!("Zitadel: http://{}:{}", ZITADEL_HOST, HTTP_PORT);
    info!("Run again to verify cached session works.");
    info!("cargo run -p example-harmony-sso -- --cleanup  # teardown");

    Ok(())
}