Skip to content

Commit e879464

Browse files
tobikristobikris
committed
feat: add support for openbao flavor
1 parent 2c75ad3 commit e879464

21 files changed

+539
-199
lines changed

‎src/exec.rs‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use secrecy::{ExposeSecret, Secret};
66
use std::collections::HashMap;
77
use tokio::io::AsyncWriteExt;
88

9-
use crate::{list_vault_pods, LABEL_KEY_VAULT_ACTIVE, LABEL_KEY_VAULT_SEALED};
9+
use crate::{list_vault_pods, Flavor};
1010

1111
#[derive(ValueEnum, Copy, Clone, Debug, PartialEq, Eq)]
1212
pub enum ExecIn {
@@ -25,11 +25,11 @@ impl std::fmt::Display for ExecIn {
2525
}
2626

2727
impl ExecIn {
28-
pub fn to_label_selector(&self) -> String {
28+
pub fn to_label_selector(&self, flavor: &str) -> String {
2929
match self {
30-
ExecIn::Active => format!("{}=true", LABEL_KEY_VAULT_ACTIVE),
31-
ExecIn::Standby => format!("{}=false", LABEL_KEY_VAULT_ACTIVE),
32-
ExecIn::Sealed => format!("{}=true", LABEL_KEY_VAULT_SEALED),
30+
ExecIn::Active => format!("{}=true", &format!("{}-active", flavor)),
31+
ExecIn::Standby => format!("{}=false", &format!("{}-active", flavor)),
32+
ExecIn::Sealed => format!("{}=true", &format!("{}-sealed", flavor)),
3333
}
3434
}
3535
}
@@ -39,10 +39,14 @@ pub async fn exec(
3939
api: &Api<Pod>,
4040
cmd: String,
4141
exec_in: ExecIn,
42+
flavor: Flavor,
4243
env: HashMap<String, Secret<String>>,
4344
) -> anyhow::Result<()> {
4445
let pods = api
45-
.list(&list_vault_pods().labels(&exec_in.to_label_selector()))
46+
.list(
47+
&list_vault_pods(&flavor.to_string())
48+
.labels(&exec_in.to_label_selector(&flavor.to_string())),
49+
)
4650
.await?;
4751
let pod = pods
4852
.items

‎src/helpers.rs‎

Lines changed: 53 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,40 +4,37 @@ use tokio::io::{AsyncRead, AsyncWrite};
44

55
use crate::{BytesBody, HttpForwarderService};
66

7-
pub const LABEL_KEY_VAULT_ACTIVE: &str = "vault-active";
8-
pub const LABEL_KEY_VAULT_SEALED: &str = "vault-sealed";
9-
10-
pub fn list_vault_pods() -> ListParams {
11-
ListParams::default().labels("app.kubernetes.io/name=vault")
7+
pub fn list_vault_pods(flavor: &str) -> ListParams {
8+
ListParams::default().labels(&format!("app.kubernetes.io/name={}", flavor))
129
}
1310

1411
/// Check if the vault pod is sealed based on its labels
1512
/// Returns an error if the pod does not have the expected labels
16-
pub fn is_sealed(pod: &Pod) -> anyhow::Result<bool> {
13+
pub fn is_sealed(pod: &Pod, flavor: &str) -> anyhow::Result<bool> {
1714
match pod.metadata.labels.as_ref() {
1815
None => Err(anyhow::anyhow!("pod does not have labels")),
19-
Some(labels) => match labels.get(LABEL_KEY_VAULT_SEALED) {
16+
Some(labels) => match labels.get(&format!("{}-sealed", flavor)) {
2017
Some(x) if x.as_str() == "true" => Ok(true),
2118
Some(x) if x.as_str() == "false" => Ok(false),
2219
_ => Err(anyhow::anyhow!(
2320
"pod does not have a {} label",
24-
LABEL_KEY_VAULT_SEALED
21+
&format!("{}-sealed", flavor)
2522
)),
2623
},
2724
}
2825
}
2926

3027
/// Check if the vault pod is active based on its labels
3128
/// Returns an error if the pod does not have the expected labels
32-
pub fn is_active(pod: &Pod) -> anyhow::Result<bool> {
29+
pub fn is_active(pod: &Pod, flavor: &str) -> anyhow::Result<bool> {
3330
match pod.metadata.labels.as_ref() {
3431
None => Err(anyhow::anyhow!("pod does not have labels")),
35-
Some(labels) => match labels.get(LABEL_KEY_VAULT_ACTIVE) {
32+
Some(labels) => match labels.get(&format!("{}-active", flavor)) {
3633
Some(x) if x.as_str() == "true" => Ok(true),
3734
Some(x) if x.as_str() == "false" => Ok(false),
3835
_ => Err(anyhow::anyhow!(
3936
"pod does not have a {} label",
40-
LABEL_KEY_VAULT_ACTIVE
37+
&format!("{}-active", flavor)
4138
)),
4239
},
4340
}
@@ -49,11 +46,50 @@ pub struct PodApi {
4946
pub api: Api<Pod>,
5047
tls: bool,
5148
domain: String,
49+
pub flavor: Flavor,
50+
}
51+
52+
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
53+
pub enum Flavor {
54+
OpenBao,
55+
Vault,
56+
}
57+
58+
impl std::fmt::Display for Flavor {
59+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60+
match self {
61+
Flavor::OpenBao => write!(f, "openbao"),
62+
Flavor::Vault => write!(f, "vault"),
63+
}
64+
}
65+
}
66+
67+
impl std::str::FromStr for Flavor {
68+
type Err = anyhow::Error;
69+
70+
fn from_str(s: &str) -> Result<Self, Self::Err> {
71+
match s.to_lowercase().as_str() {
72+
"openbao" => Ok(Flavor::OpenBao),
73+
"vault" => Ok(Flavor::Vault),
74+
_ => Err(anyhow::anyhow!("invalid flavor: {}", s)),
75+
}
76+
}
77+
}
78+
79+
impl Flavor {
80+
pub fn container_name(&self) -> String {
81+
self.to_string()
82+
}
5283
}
5384

5485
impl PodApi {
55-
pub fn new(api: Api<Pod>, tls: bool, domain: String) -> Self {
56-
Self { api, tls, domain }
86+
pub fn new(api: Api<Pod>, tls: bool, domain: String, flavor: Flavor) -> Self {
87+
Self {
88+
api,
89+
tls,
90+
domain,
91+
flavor,
92+
}
5793
}
5894
}
5995

@@ -91,10 +127,11 @@ impl PodApi {
91127
/// Wrapper around the kube::Api type for the Vault statefulset
92128
pub struct StatefulSetApi {
93129
pub api: Api<StatefulSet>,
130+
pub flavor: Flavor,
94131
}
95132

96-
impl From<Api<StatefulSet>> for StatefulSetApi {
97-
fn from(api: Api<StatefulSet>) -> Self {
98-
Self { api }
133+
impl StatefulSetApi {
134+
pub fn new(api: Api<StatefulSet>, flavor: Flavor) -> Self {
135+
Self { api, flavor }
99136
}
100137
}

‎src/init.rs‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use kube::Api;
66
use secrecy::Secret;
77
use tracing::*;
88

9-
use crate::{init_request, raft_join_request, BytesBody, HttpRequest, PodApi, VAULT_PORT};
9+
use crate::{init_request, raft_join_request, BytesBody, Flavor, HttpRequest, PodApi, VAULT_PORT};
1010

1111
#[derive(Debug, serde::Serialize)]
1212
pub struct InitRequest {
@@ -105,12 +105,17 @@ where
105105
}
106106

107107
#[tracing::instrument(skip_all)]
108-
pub async fn init(domain: String, api: &Api<Pod>, pod_name: &str) -> anyhow::Result<InitResult> {
108+
pub async fn init(
109+
domain: String,
110+
api: &Api<Pod>,
111+
flavor: Flavor,
112+
pod_name: &str,
113+
) -> anyhow::Result<InitResult> {
109114
let pod = api.get(pod_name).await?;
110115

111116
info!("initializing: {}", pod_name);
112117

113-
let pods = PodApi::new(api.clone(), true, domain);
118+
let pods = PodApi::new(api.clone(), true, domain, flavor);
114119
let mut pf = pods
115120
.http(
116121
pod.metadata
@@ -158,6 +163,7 @@ pub async fn init(domain: String, api: &Api<Pod>, pod_name: &str) -> anyhow::Res
158163
pub async fn raft_join(
159164
domain: String,
160165
api: &Api<Pod>,
166+
flavor: Flavor,
161167
pod_name: &str,
162168
join_to: &str,
163169
) -> anyhow::Result<()> {
@@ -172,7 +178,7 @@ pub async fn raft_join(
172178
join_to,
173179
);
174180

175-
let pods = PodApi::new(api.clone(), true, domain);
181+
let pods = PodApi::new(api.clone(), true, domain, flavor);
176182
let mut pf = pods
177183
.http(
178184
pod.metadata

‎src/main.rs‎

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ use std::io;
1010
use std::str::FromStr;
1111
use tokio::task::spawn_blocking;
1212
use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
13+
use vault_mgmt_lib::Flavor;
1314

1415
use vault_mgmt_lib::{
1516
construct_table, is_statefulset_ready, GetUnsealKeys, GetUnsealKeysFromVault, StepDown,
@@ -48,6 +49,17 @@ struct Cli {
4849
#[arg(long)]
4950
no_tls: bool,
5051

52+
/// Server Flavor
53+
#[arg(
54+
short = 'f',
55+
long,
56+
default_value = "vault",
57+
value_parser = clap::builder::PossibleValuesParser::new(
58+
["openbao", "vault"]
59+
).map(|s| Flavor::from_str(&s).expect("invalid flavor")),
60+
)]
61+
flavor: Flavor,
62+
5163
/// Subcommand to run
5264
#[command(subcommand)]
5365
command: Commands,
@@ -188,7 +200,7 @@ async fn main() -> anyhow::Result<()> {
188200
}
189201
Commands::Show {} => {
190202
let api = setup_api(&cli.namespace).await?;
191-
let table = construct_table(&api).await?;
203+
let table = construct_table(&api, cli.flavor).await?;
192204

193205
table.printstd();
194206
}
@@ -200,18 +212,21 @@ async fn main() -> anyhow::Result<()> {
200212
} => {
201213
let api = setup_api(&cli.namespace).await?;
202214
let env = collect_env(env, env_keys)?;
203-
exec(&api, cmd.join(" "), exec_in, env).await?;
215+
exec(&api, cmd.join(" "), exec_in, cli.flavor, env).await?;
204216
}
205217
Commands::StepDown { token } => {
206218
let api = setup_api(&cli.namespace).await?;
207219
let active = api
208-
.list(&list_vault_pods().labels(&ExecIn::Active.to_label_selector()))
220+
.list(
221+
&list_vault_pods(&cli.flavor.to_string())
222+
.labels(&ExecIn::Active.to_label_selector(&cli.flavor.to_string())),
223+
)
209224
.await?;
210225
let active = active.iter().next().ok_or(anyhow::anyhow!(
211226
"no active vault pod found. is vault sealed?"
212227
))?;
213228

214-
PodApi::new(api, !cli.no_tls, cli.domain)
229+
PodApi::new(api, !cli.no_tls, cli.domain, cli.flavor)
215230
.http(
216231
active
217232
.metadata
@@ -240,7 +255,7 @@ async fn main() -> anyhow::Result<()> {
240255
key_cmd,
241256
} => {
242257
let api = setup_api(&cli.namespace).await?;
243-
let sealed = list_sealed_pods(&api).await?;
258+
let sealed = list_sealed_pods(&api, &cli.flavor.to_string()).await?;
244259

245260
if sealed.is_empty() {
246261
return Ok(());
@@ -277,7 +292,7 @@ async fn main() -> anyhow::Result<()> {
277292
}
278293

279294
for pod in sealed.iter() {
280-
PodApi::new(api.clone(), !cli.no_tls, cli.domain.clone())
295+
PodApi::new(api.clone(), !cli.no_tls, cli.domain.clone(), cli.flavor)
281296
.http(
282297
pod.metadata
283298
.name
@@ -333,10 +348,10 @@ async fn main() -> anyhow::Result<()> {
333348

334349
let sts = stss.get(&cli.statefulset).await?;
335350

336-
StatefulSetApi::from(stss.clone())
351+
StatefulSetApi::new(stss.clone(), cli.flavor)
337352
.upgrade(
338353
sts.clone(),
339-
&PodApi::new(pods.clone(), !cli.no_tls, cli.domain),
354+
&PodApi::new(pods.clone(), !cli.no_tls, cli.domain, cli.flavor),
340355
token,
341356
!do_not_unseal,
342357
force_upgrade,

‎src/show.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@ use k8s_openapi::api::core::v1::Pod;
22
use kube::api::Api;
33
use prettytable::{color, Attr, Cell, Row, Table};
44

5-
use crate::list_vault_pods;
5+
use crate::{list_vault_pods, Flavor};
66

77
#[tracing::instrument(skip_all)]
8-
pub async fn construct_table(api: &Api<Pod>) -> anyhow::Result<Table> {
8+
pub async fn construct_table(api: &Api<Pod>, flavor: Flavor) -> anyhow::Result<Table> {
99
let mut table = Table::new();
1010
table.set_titles(row![
1111
"NAME",
@@ -17,7 +17,7 @@ pub async fn construct_table(api: &Api<Pod>) -> anyhow::Result<Table> {
1717
"READY",
1818
]);
1919

20-
let pods = api.list(&list_vault_pods()).await?;
20+
let pods = api.list(&list_vault_pods(&flavor.to_string())).await?;
2121

2222
let get_vault_label = |pod: &Pod, label: &str| match pod.metadata.labels {
2323
Some(ref labels) => labels
@@ -51,22 +51,22 @@ pub async fn construct_table(api: &Api<Pod>) -> anyhow::Result<Table> {
5151
.clone()
5252
.ok_or(anyhow::anyhow!("container does not have an image"))?;
5353

54-
let initialized = get_vault_label(p, "vault-initialized");
54+
let initialized = get_vault_label(p, &format!("{}-initialized", flavor));
5555
let initialized =
5656
Cell::new(&initialized).with_style(Attr::ForegroundColor(match initialized.as_str() {
5757
"true" => color::GREEN,
5858
"false" => color::RED,
5959
_ => color::YELLOW,
6060
}));
6161

62-
let sealed = get_vault_label(p, "vault-sealed");
62+
let sealed = get_vault_label(p, &format!("{}-sealed", flavor));
6363
let sealed = Cell::new(&sealed).with_style(Attr::ForegroundColor(match sealed.as_str() {
6464
"true" => color::RED,
6565
"false" => color::GREEN,
6666
_ => color::YELLOW,
6767
}));
6868

69-
let active = get_vault_label(p, "vault-active");
69+
let active = get_vault_label(p, &format!("{}-active", flavor));
7070
let active = Cell::new(&active).with_style(Attr::ForegroundColor(match active.as_str() {
7171
"true" => color::GREEN,
7272
"false" => color::WHITE,

‎src/unseal.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,9 @@ pub async fn get_unseal_keys(key_cmd: &str) -> anyhow::Result<Vec<Secret<String>
2828
}
2929

3030
/// List all pods that are sealed
31-
pub async fn list_sealed_pods(api: &Api<Pod>) -> anyhow::Result<Vec<Pod>> {
31+
pub async fn list_sealed_pods(api: &Api<Pod>, flavor: &str) -> anyhow::Result<Vec<Pod>> {
3232
let pods = api
33-
.list(&list_vault_pods().labels(&ExecIn::Sealed.to_label_selector()))
33+
.list(&list_vault_pods(flavor).labels(&ExecIn::Sealed.to_label_selector(flavor)))
3434
.await?;
3535

3636
Ok(pods.items)
@@ -295,7 +295,7 @@ mod tests {
295295
async fn get_sealed_pods_returns_sealed_pods() {
296296
let (api, service, cancel) = setup().await;
297297

298-
let pods = list_sealed_pods(&api).await.unwrap();
298+
let pods = list_sealed_pods(&api, "vault").await.unwrap();
299299

300300
assert_eq!(pods.len(), 3);
301301

0 commit comments

Comments
 (0)