feat: add support for openbao flavor

Author: tobikris

src/exec.rs

@@ -6,7 +6,7 @@ use secrecy::{ExposeSecret, Secret};
 use std::collections::HashMap;
 use tokio::io::AsyncWriteExt;
 
-use crate::{list_vault_pods, LABEL_KEY_VAULT_ACTIVE, LABEL_KEY_VAULT_SEALED};
+use crate::{list_vault_pods, Flavor};
 
 #[derive(ValueEnum, Copy, Clone, Debug, PartialEq, Eq)]
 pub enum ExecIn {
@@ -25,11 +25,11 @@ impl std::fmt::Display for ExecIn {
 }
 
 impl ExecIn {
-    pub fn to_label_selector(&self) -> String {
+    pub fn to_label_selector(&self, flavor: &str) -> String {
         match self {
-            ExecIn::Active => format!("{}=true", LABEL_KEY_VAULT_ACTIVE),
-            ExecIn::Standby => format!("{}=false", LABEL_KEY_VAULT_ACTIVE),
-            ExecIn::Sealed => format!("{}=true", LABEL_KEY_VAULT_SEALED),
+            ExecIn::Active => format!("{}=true", &format!("{}-active", flavor)),
+            ExecIn::Standby => format!("{}=false", &format!("{}-active", flavor)),
+            ExecIn::Sealed => format!("{}=true", &format!("{}-sealed", flavor)),
         }
     }
 }
@@ -39,10 +39,14 @@ pub async fn exec(
     api: &Api<Pod>,
     cmd: String,
     exec_in: ExecIn,
+    flavor: Flavor,
     env: HashMap<String, Secret<String>>,
 ) -> anyhow::Result<()> {
     let pods = api
-        .list(&list_vault_pods().labels(&exec_in.to_label_selector()))
+        .list(
+            &list_vault_pods(&flavor.to_string())
+                .labels(&exec_in.to_label_selector(&flavor.to_string())),
+        )
         .await?;
     let pod = pods
         .items

src/helpers.rs

@@ -4,40 +4,37 @@ use tokio::io::{AsyncRead, AsyncWrite};
 
 use crate::{BytesBody, HttpForwarderService};
 
-pub const LABEL_KEY_VAULT_ACTIVE: &str = "vault-active";
-pub const LABEL_KEY_VAULT_SEALED: &str = "vault-sealed";
-
-pub fn list_vault_pods() -> ListParams {
-    ListParams::default().labels("app.kubernetes.io/name=vault")
+pub fn list_vault_pods(flavor: &str) -> ListParams {
+    ListParams::default().labels(&format!("app.kubernetes.io/name={}", flavor))
 }
 
 /// Check if the vault pod is sealed based on its labels
 /// Returns an error if the pod does not have the expected labels
-pub fn is_sealed(pod: &Pod) -> anyhow::Result<bool> {
+pub fn is_sealed(pod: &Pod, flavor: &str) -> anyhow::Result<bool> {
     match pod.metadata.labels.as_ref() {
         None => Err(anyhow::anyhow!("pod does not have labels")),
-        Some(labels) => match labels.get(LABEL_KEY_VAULT_SEALED) {
+        Some(labels) => match labels.get(&format!("{}-sealed", flavor)) {
             Some(x) if x.as_str() == "true" => Ok(true),
             Some(x) if x.as_str() == "false" => Ok(false),
             _ => Err(anyhow::anyhow!(
                 "pod does not have a {} label",
-                LABEL_KEY_VAULT_SEALED
+                &format!("{}-sealed", flavor)
             )),
         },
     }
 }
 
 /// Check if the vault pod is active based on its labels
 /// Returns an error if the pod does not have the expected labels
-pub fn is_active(pod: &Pod) -> anyhow::Result<bool> {
+pub fn is_active(pod: &Pod, flavor: &str) -> anyhow::Result<bool> {
     match pod.metadata.labels.as_ref() {
         None => Err(anyhow::anyhow!("pod does not have labels")),
-        Some(labels) => match labels.get(LABEL_KEY_VAULT_ACTIVE) {
+        Some(labels) => match labels.get(&format!("{}-active", flavor)) {
             Some(x) if x.as_str() == "true" => Ok(true),
             Some(x) if x.as_str() == "false" => Ok(false),
             _ => Err(anyhow::anyhow!(
                 "pod does not have a {} label",
-                LABEL_KEY_VAULT_ACTIVE
+                &format!("{}-active", flavor)
             )),
         },
     }
@@ -49,11 +46,50 @@ pub struct PodApi {
     pub api: Api<Pod>,
     tls: bool,
     domain: String,
+    pub flavor: Flavor,
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum Flavor {
+    OpenBao,
+    Vault,
+}
+
+impl std::fmt::Display for Flavor {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Flavor::OpenBao => write!(f, "openbao"),
+            Flavor::Vault => write!(f, "vault"),
+        }
+    }
+}
+
+impl std::str::FromStr for Flavor {
+    type Err = anyhow::Error;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_lowercase().as_str() {
+            "openbao" => Ok(Flavor::OpenBao),
+            "vault" => Ok(Flavor::Vault),
+            _ => Err(anyhow::anyhow!("invalid flavor: {}", s)),
+        }
+    }
+}
+
+impl Flavor {
+    pub fn container_name(&self) -> String {
+        self.to_string()
+    }
 }
 
 impl PodApi {
-    pub fn new(api: Api<Pod>, tls: bool, domain: String) -> Self {
-        Self { api, tls, domain }
+    pub fn new(api: Api<Pod>, tls: bool, domain: String, flavor: Flavor) -> Self {
+        Self {
+            api,
+            tls,
+            domain,
+            flavor,
+        }
     }
 }
 
@@ -91,10 +127,11 @@ impl PodApi {
 /// Wrapper around the kube::Api type for the Vault statefulset
 pub struct StatefulSetApi {
     pub api: Api<StatefulSet>,
+    pub flavor: Flavor,
 }
 
-impl From<Api<StatefulSet>> for StatefulSetApi {
-    fn from(api: Api<StatefulSet>) -> Self {
-        Self { api }
+impl StatefulSetApi {
+    pub fn new(api: Api<StatefulSet>, flavor: Flavor) -> Self {
+        Self { api, flavor }
     }
 }

src/init.rs

@@ -6,7 +6,7 @@ use kube::Api;
 use secrecy::Secret;
 use tracing::*;
 
-use crate::{init_request, raft_join_request, BytesBody, HttpRequest, PodApi, VAULT_PORT};
+use crate::{init_request, raft_join_request, BytesBody, Flavor, HttpRequest, PodApi, VAULT_PORT};
 
 #[derive(Debug, serde::Serialize)]
 pub struct InitRequest {
@@ -105,12 +105,17 @@ where
 }
 
 #[tracing::instrument(skip_all)]
-pub async fn init(domain: String, api: &Api<Pod>, pod_name: &str) -> anyhow::Result<InitResult> {
+pub async fn init(
+    domain: String,
+    api: &Api<Pod>,
+    flavor: Flavor,
+    pod_name: &str,
+) -> anyhow::Result<InitResult> {
     let pod = api.get(pod_name).await?;
 
     info!("initializing: {}", pod_name);
 
-    let pods = PodApi::new(api.clone(), true, domain);
+    let pods = PodApi::new(api.clone(), true, domain, flavor);
     let mut pf = pods
         .http(
             pod.metadata
@@ -158,6 +163,7 @@ pub async fn init(domain: String, api: &Api<Pod>, pod_name: &str) -> anyhow::Res
 pub async fn raft_join(
     domain: String,
     api: &Api<Pod>,
+    flavor: Flavor,
     pod_name: &str,
     join_to: &str,
 ) -> anyhow::Result<()> {
@@ -172,7 +178,7 @@ pub async fn raft_join(
         join_to,
     );
 
-    let pods = PodApi::new(api.clone(), true, domain);
+    let pods = PodApi::new(api.clone(), true, domain, flavor);
     let mut pf = pods
         .http(
             pod.metadata

src/main.rs

@@ -10,6 +10,7 @@ use std::io;
 use std::str::FromStr;
 use tokio::task::spawn_blocking;
 use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
+use vault_mgmt_lib::Flavor;
 
 use vault_mgmt_lib::{
     construct_table, is_statefulset_ready, GetUnsealKeys, GetUnsealKeysFromVault, StepDown,
@@ -48,6 +49,17 @@ struct Cli {
     #[arg(long)]
     no_tls: bool,
 
+    /// Server Flavor
+    #[arg(
+        short = 'f',
+        long,
+        default_value = "vault",
+        value_parser = clap::builder::PossibleValuesParser::new(
+            ["openbao", "vault"]
+        ).map(|s| Flavor::from_str(&s).expect("invalid flavor")),
+    )]
+    flavor: Flavor,
+
     /// Subcommand to run
     #[command(subcommand)]
     command: Commands,
@@ -188,7 +200,7 @@ async fn main() -> anyhow::Result<()> {
         }
         Commands::Show {} => {
             let api = setup_api(&cli.namespace).await?;
-            let table = construct_table(&api).await?;
+            let table = construct_table(&api, cli.flavor).await?;
 
             table.printstd();
         }
@@ -200,18 +212,21 @@ async fn main() -> anyhow::Result<()> {
         } => {
             let api = setup_api(&cli.namespace).await?;
             let env = collect_env(env, env_keys)?;
-            exec(&api, cmd.join(" "), exec_in, env).await?;
+            exec(&api, cmd.join(" "), exec_in, cli.flavor, env).await?;
         }
         Commands::StepDown { token } => {
             let api = setup_api(&cli.namespace).await?;
             let active = api
-                .list(&list_vault_pods().labels(&ExecIn::Active.to_label_selector()))
+                .list(
+                    &list_vault_pods(&cli.flavor.to_string())
+                        .labels(&ExecIn::Active.to_label_selector(&cli.flavor.to_string())),
+                )
                 .await?;
             let active = active.iter().next().ok_or(anyhow::anyhow!(
                 "no active vault pod found. is vault sealed?"
             ))?;
 
-            PodApi::new(api, !cli.no_tls, cli.domain)
+            PodApi::new(api, !cli.no_tls, cli.domain, cli.flavor)
                 .http(
                     active
                         .metadata
@@ -240,7 +255,7 @@ async fn main() -> anyhow::Result<()> {
             key_cmd,
         } => {
             let api = setup_api(&cli.namespace).await?;
-            let sealed = list_sealed_pods(&api).await?;
+            let sealed = list_sealed_pods(&api, &cli.flavor.to_string()).await?;
 
             if sealed.is_empty() {
                 return Ok(());
@@ -277,7 +292,7 @@ async fn main() -> anyhow::Result<()> {
             }
 
             for pod in sealed.iter() {
-                PodApi::new(api.clone(), !cli.no_tls, cli.domain.clone())
+                PodApi::new(api.clone(), !cli.no_tls, cli.domain.clone(), cli.flavor)
                     .http(
                         pod.metadata
                             .name
@@ -333,10 +348,10 @@ async fn main() -> anyhow::Result<()> {
 
             let sts = stss.get(&cli.statefulset).await?;
 
-            StatefulSetApi::from(stss.clone())
+            StatefulSetApi::new(stss.clone(), cli.flavor)
                 .upgrade(
                     sts.clone(),
-                    &PodApi::new(pods.clone(), !cli.no_tls, cli.domain),
+                    &PodApi::new(pods.clone(), !cli.no_tls, cli.domain, cli.flavor),
                     token,
                     !do_not_unseal,
                     force_upgrade,

src/show.rs

@@ -2,10 +2,10 @@ use k8s_openapi::api::core::v1::Pod;
 use kube::api::Api;
 use prettytable::{color, Attr, Cell, Row, Table};
 
-use crate::list_vault_pods;
+use crate::{list_vault_pods, Flavor};
 
 #[tracing::instrument(skip_all)]
-pub async fn construct_table(api: &Api<Pod>) -> anyhow::Result<Table> {
+pub async fn construct_table(api: &Api<Pod>, flavor: Flavor) -> anyhow::Result<Table> {
     let mut table = Table::new();
     table.set_titles(row![
         "NAME",
@@ -17,7 +17,7 @@ pub async fn construct_table(api: &Api<Pod>) -> anyhow::Result<Table> {
         "READY",
     ]);
 
-    let pods = api.list(&list_vault_pods()).await?;
+    let pods = api.list(&list_vault_pods(&flavor.to_string())).await?;
 
     let get_vault_label = |pod: &Pod, label: &str| match pod.metadata.labels {
         Some(ref labels) => labels
@@ -51,22 +51,22 @@ pub async fn construct_table(api: &Api<Pod>) -> anyhow::Result<Table> {
             .clone()
             .ok_or(anyhow::anyhow!("container does not have an image"))?;
 
-        let initialized = get_vault_label(p, "vault-initialized");
+        let initialized = get_vault_label(p, &format!("{}-initialized", flavor.to_string()));
         let initialized =
             Cell::new(&initialized).with_style(Attr::ForegroundColor(match initialized.as_str() {
                 "true" => color::GREEN,
                 "false" => color::RED,
                 _ => color::YELLOW,
             }));
 
-        let sealed = get_vault_label(p, "vault-sealed");
+        let sealed = get_vault_label(p, &format!("{}-sealed", flavor.to_string()));
         let sealed = Cell::new(&sealed).with_style(Attr::ForegroundColor(match sealed.as_str() {
             "true" => color::RED,
             "false" => color::GREEN,
             _ => color::YELLOW,
         }));
 
-        let active = get_vault_label(p, "vault-active");
+        let active = get_vault_label(p, &format!("{}-active", flavor.to_string()));
         let active = Cell::new(&active).with_style(Attr::ForegroundColor(match active.as_str() {
             "true" => color::GREEN,
             "false" => color::WHITE,

src/unseal.rs

@@ -28,9 +28,9 @@ pub async fn get_unseal_keys(key_cmd: &str) -> anyhow::Result<Vec<Secret<String>
 }
 
 /// List all pods that are sealed
-pub async fn list_sealed_pods(api: &Api<Pod>) -> anyhow::Result<Vec<Pod>> {
+pub async fn list_sealed_pods(api: &Api<Pod>, flavor: &str) -> anyhow::Result<Vec<Pod>> {
     let pods = api
-        .list(&list_vault_pods().labels(&ExecIn::Sealed.to_label_selector()))
+        .list(&list_vault_pods(flavor).labels(&ExecIn::Sealed.to_label_selector(flavor)))
         .await?;
 
     Ok(pods.items)
@@ -295,7 +295,7 @@ mod tests {
     async fn get_sealed_pods_returns_sealed_pods() {
         let (api, service, cancel) = setup().await;
 
-        let pods = list_sealed_pods(&api).await.unwrap();
+        let pods = list_sealed_pods(&api, "vault").await.unwrap();
 
         assert_eq!(pods.len(), 3);