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.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"

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<String>,
     region: Option<String>,
@@ -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<String>) -> 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<StackStatus> = 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<StackStatus> = 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(())

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!(),
                 }
             }
         };

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,