From 9a1ab8f31c57e6a71ba40337b83a75d261d92e65 Mon Sep 17 00:00:00 2001 From: Akseli Lukkarila Date: Thu, 17 Oct 2024 15:22:42 +0300 Subject: [PATCH] Rust: wait for stack update in CLI (#538) * rust: wait for stack update to finish like for init * rust: add id command for checking aws account id * rust: bump version * use cf client from vault --- rust/Cargo.lock | 2 +- rust/Cargo.toml | 2 +- rust/src/cli.rs | 83 ++++++++++++++++++++++++++++++++++++++++++----- rust/src/main.rs | 13 ++++++-- rust/src/vault.rs | 6 ++-- 5 files changed, 90 insertions(+), 16 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 29b971c0..bf81f1d3 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -1290,7 +1290,7 @@ dependencies = [ [[package]] name = "nitor-vault" -version = "1.2.2" +version = "1.3.0" dependencies = [ "aes-gcm", "anyhow", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 90baa16a..1bf33391 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "nitor-vault" -version = "1.2.2" +version = "1.3.0" edition = "2021" description = "Encrypted AWS key-value storage utility." license = "Apache-2.0" diff --git a/rust/src/cli.rs b/rust/src/cli.rs index 227c54fe..74863db0 100644 --- a/rust/src/cli.rs +++ b/rust/src/cli.rs @@ -8,8 +8,11 @@ use tokio::time::Duration; use nitor_vault::{cloudformation, CreateStackResult, Value, Vault}; -const INIT_WAIT_ANIMATION_DURATION: Duration = Duration::from_millis(600); +static WAIT_ANIMATION_DURATION: Duration = Duration::from_millis(500); +static CLEAR_LINE: &str = "\x1b[2K"; +static WAIT_DOTS: [&str; 4] = [".", "..", "...", ""]; +/// Initialize a new vault stack with Cloudformation and wait for creation to finish. pub async fn init_vault_stack( stack_name: Option, region: Option, @@ -39,6 +42,16 @@ pub async fn init_vault_stack( Ok(()) } +/// Update existing Cloudformation vault stack and wait for update to finish. +pub async fn update_vault_stack(vault: &Vault) -> Result<()> { + vault + .update_stack() + .await + .with_context(|| "Failed to update vault stack".red())?; + + wait_for_stack_update_to_finish(vault).await +} + /// Store a key-value pair pub async fn store( vault: &Vault, @@ -145,28 +158,40 @@ pub async fn exists(vault: &Vault, key: &str) -> Result<()> { }) } +/// Print the information from AWS STS "get caller identity" call +pub async fn print_aws_account(region: Option) -> Result<()> { + let config = Vault::get_aws_config(region).await; + let client = aws_sdk_sts::Client::new(&config); + let result = client.get_caller_identity().send().await?; + println!( + "user: {}\naccount: {}\narn: {}", + result.user_id.unwrap_or_else(|| "None".to_string()), + result.account.unwrap_or_else(|| "None".to_string()), + result.arn.unwrap_or_else(|| "None".to_string()) + ); + Ok(()) +} + /// Poll Cloudformation for stack status until it has been created or creation failed. async fn wait_for_stack_creation_to_finish( config: &aws_config::SdkConfig, stack_name: &str, ) -> Result<()> { let client = aws_sdk_cloudformation::Client::new(config); - let clear_line = "\x1b[2K"; - let dots = [".", "..", "...", ""]; let mut last_status: Option = None; loop { let stack_data = cloudformation::get_stack_data(&client, stack_name).await?; if let Some(ref status) = stack_data.status { match status { StackStatus::CreateComplete => { - println!("{clear_line}{stack_data}"); + println!("{CLEAR_LINE}{stack_data}"); println!("{}", "Stack creation completed successfully".green()); break; } StackStatus::CreateFailed | StackStatus::RollbackFailed | StackStatus::RollbackComplete => { - println!("{clear_line}{stack_data}"); + println!("{CLEAR_LINE}{stack_data}"); anyhow::bail!("Stack creation failed"); } _ => { @@ -176,13 +201,55 @@ async fn wait_for_stack_creation_to_finish( println!("status: {status}"); } // Continue waiting for stack creation to complete - for dot in &dots { - print!("\r{clear_line}{dot}"); + for dot in WAIT_DOTS { + print!("\r{CLEAR_LINE}{dot}"); std::io::stdout().flush()?; - tokio::time::sleep(INIT_WAIT_ANIMATION_DURATION).await; + tokio::time::sleep(WAIT_ANIMATION_DURATION).await; } } } + } else { + anyhow::bail!("Failed to get stack status for stack '{stack_name}'"); + } + } + Ok(()) +} + +/// Poll Cloudformation for stack status until it has been updated or update failed. +async fn wait_for_stack_update_to_finish(vault: &Vault) -> Result<()> { + let mut last_status: Option = None; + loop { + let stack_data = vault.stack_status().await?; + if let Some(ref status) = stack_data.status { + match status { + StackStatus::UpdateComplete => { + println!("{CLEAR_LINE}{stack_data}"); + println!("{}", "Stack update completed successfully".green()); + break; + } + StackStatus::UpdateFailed | StackStatus::RollbackFailed => { + println!("{CLEAR_LINE}{stack_data}"); + anyhow::bail!("Stack update failed"); + } + _ => { + // Print status if it has changed + if last_status.as_ref() != Some(status) { + last_status = Some(status.clone()); + println!("status: {status}"); + } + // Continue waiting for stack update to complete + for dot in WAIT_DOTS { + print!("\r{CLEAR_LINE}{dot}"); + std::io::stdout().flush()?; + tokio::time::sleep(WAIT_ANIMATION_DURATION).await; + } + } + } + } else { + anyhow::bail!( + "Failed to get stack status for stack '{}'", + vault.cloudformation_params.stack_name + ); } } Ok(()) diff --git a/rust/src/main.rs b/rust/src/main.rs index ccd48573..3a682a40 100644 --- a/rust/src/main.rs +++ b/rust/src/main.rs @@ -65,6 +65,10 @@ pub enum Command { #[command(long_flag("info"))] Info {}, + /// Print AWS user account information + #[command(long_flag("id"))] + Id {}, + /// Print vault stack information #[command(long_flag("status"))] Status {}, @@ -162,6 +166,7 @@ pub enum Command { }, } +#[allow(clippy::match_same_arms)] #[tokio::main] async fn main() -> Result<()> { let args = Args::parse(); @@ -183,11 +188,13 @@ async fn main() -> Result<()> { ) .await .with_context(|| "Failed to create vault from given params".red())?; - vault - .update_stack() + cli::update_vault_stack(&vault) .await .with_context(|| "Failed to update vault stack".red())?; } + Command::Id {} => { + cli::print_aws_account(args.region).await?; + } Command::All {} | Command::Delete { .. } | Command::Describe {} @@ -225,9 +232,9 @@ async fn main() -> Result<()> { } => cli::store(&vault, key, value, file, value_argument, overwrite).await?, // These are here again instead of a `_` so that if new commands are added, // there is an error about missing handling for that. - #[allow(clippy::match_same_arms)] Command::Init { .. } => unreachable!(), Command::Update { .. } => unreachable!(), + Command::Id { .. } => unreachable!(), } } }; diff --git a/rust/src/vault.rs b/rust/src/vault.rs index 28cf3ad6..05042d5b 100644 --- a/rust/src/vault.rs +++ b/rust/src/vault.rs @@ -29,10 +29,10 @@ use crate::{CreateStackResult, EncryptObject, Meta, S3DataKeys}; pub struct Vault { /// AWS region to use with Vault. /// Will fall back to default provider if nothing is specified. - region: Region, + pub region: Region, /// Prefix for key name - prefix: String, - cloudformation_params: CloudFormationParams, + pub prefix: String, + pub cloudformation_params: CloudFormationParams, cf: CloudFormationClient, kms: KmsClient, s3: S3Client,