use std::cmp::Reverse; use std::fmt; use std::io::{self, IsTerminal as _}; use anyhow::{anyhow, Result}; use chrono::{DateTime, Utc}; use comfy_table::Cell; use inquire::{Confirm, InquireError, Select}; use serde::Serialize; use warp_cli::agent::OutputFormat; use warp_cli::api_key::{ ApiKeyCommand, ApiKeyExpirationArgs, ApiKeySortByArg, CreateApiKeyArgs, ExpireApiKeyArgs, ListApiKeysArgs, }; use warp_cli::{GlobalOptions, SortOrderArg}; use warp_graphql::mutations::expire_api_key::ExpireApiKeyResult; use warp_graphql::mutations::generate_api_key::GenerateApiKeyResult; use warp_graphql::queries::api_keys::ApiKeyProperties; use warp_graphql::scalars::Time; use warpui::platform::TerminationMode; use warpui::{AppContext, ModelContext, SingletonEntity}; use super::output::{self, TableFormat}; use crate::server::ids::ApiKeyUid; use crate::util::time_format::format_approx_duration_from_now_utc; use crate::ServerApiProvider; /// Run API key-related commands. pub fn run( ctx: &mut AppContext, global_options: GlobalOptions, command: ApiKeyCommand, ) -> Result<()> { let runner = ctx.add_singleton_model(|_ctx| ApiKeyCommandRunner); match command { ApiKeyCommand::List(args) => { runner.update(ctx, |runner, ctx| { runner.list(global_options.output_format, args, ctx) }); Ok(()) } ApiKeyCommand::Create(args) => { runner.update(ctx, |runner, ctx| { runner.create(global_options.output_format, args, ctx) }); Ok(()) } ApiKeyCommand::Expire(args) => { runner.update(ctx, |runner, ctx| { runner.expire(global_options.output_format, args, ctx) }); Ok(()) } } } struct ApiKeyCommandRunner; impl ApiKeyCommandRunner { fn list( &self, output_format: OutputFormat, args: ListApiKeysArgs, ctx: &mut ModelContext, ) { let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client(); ctx.spawn( async move { let mut keys: Vec<_> = auth_client .list_api_keys() .await? .into_iter() .map(ApiKeyInfo::from) .collect(); sort_api_keys(&mut keys, args.sort_by, args.sort_order); if args.json_output.force_json_output() { output::print_raw_json(serde_json::to_value(&keys)?, &args.json_output)?; } else { output::print_list(keys, output_format); } Ok(()) }, |_, result: Result<()>, ctx| finish_command(result, ctx), ); } fn create( &self, output_format: OutputFormat, args: CreateApiKeyArgs, ctx: &mut ModelContext, ) { let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client(); ctx.spawn( async move { let json_output = args.json_output; let expires_at = expires_at_from_args(args.expiration)?; let agent_uid = args.agent_uid.map(cynic::Id::new); let result = auth_client .create_api_key(args.name, None, agent_uid, expires_at) .await?; let result = match result { GenerateApiKeyResult::GenerateApiKeyOutput(output) => CreatedApiKeyInfo { raw_api_key: output.raw_api_key, api_key: ApiKeyInfo::from(output.api_key), }, GenerateApiKeyResult::UserFacingError(e) => { return Err(anyhow!( warp_graphql::client::get_user_facing_error_message(e) )); } GenerateApiKeyResult::Unknown => { return Err(anyhow!("failed to create API key")) } }; print_created_api_key(result, output_format, json_output)?; Ok(()) }, |_, result: Result<()>, ctx| finish_command(result, ctx), ); } fn expire( &self, output_format: OutputFormat, args: ExpireApiKeyArgs, ctx: &mut ModelContext, ) { let key_identifier = args.key_uid; let force = args.force; let json_output = args.json_output; let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client(); ctx.spawn( async move { let keys = auth_client .list_api_keys() .await? .into_iter() .map(ApiKeyInfo::from) .collect(); Ok(keys) }, move |_, result: Result>, ctx| { let keys = match result { Ok(keys) => keys, Err(err) => { super::report_fatal_error(err, ctx); return; } }; let key = match resolve_api_key_identifier(&keys, &key_identifier) { Ok(Some(key)) => key, Ok(None) => { ctx.terminate_app(TerminationMode::ForceTerminate, None); return; } Err(err) => { super::report_fatal_error(err, ctx); return; } }; if !force { if !io::stdin().is_terminal() { super::report_fatal_error( anyhow!( "Refusing to expire API key without confirmation in non-interactive mode (use --force to bypass)" ), ctx, ); return; } let prompt = format!("Expire API key '{key}'?"); let should_expire = match Confirm::new(&prompt) .with_default(false) .with_help_message("This action takes effect immediately") .prompt() { Ok(should_expire) => should_expire, Err(InquireError::OperationCanceled | InquireError::OperationInterrupted) => { ctx.terminate_app(TerminationMode::ForceTerminate, None); return; } Err(err) => { super::report_fatal_error(err.into(), ctx); return; } }; if !should_expire { println!("Expiration cancelled"); ctx.terminate_app(TerminationMode::ForceTerminate, None); return; } } let uid = ApiKeyUid::from(key.uid); let auth_client = ServerApiProvider::as_ref(ctx).get_auth_client(); ctx.spawn( async move { let result = auth_client.expire_api_key(&uid).await?; let expired = match result { ExpireApiKeyResult::ExpireApiKeyOutput(output) => output.success, ExpireApiKeyResult::UserFacingError(e) => { return Err(anyhow!( warp_graphql::client::get_user_facing_error_message(e) )); } ExpireApiKeyResult::Unknown => { return Err(anyhow!("failed to expire API key")) } }; print_expire_api_key_result( uid.to_string(), expired, output_format, json_output, )?; Ok(()) }, |_, result: Result<()>, ctx| finish_command(result, ctx), ); }, ); } } impl warpui::Entity for ApiKeyCommandRunner { type Event = (); } impl SingletonEntity for ApiKeyCommandRunner {} #[derive(Debug, Clone, Serialize, PartialEq, Eq)] struct ApiKeyInfo { uid: String, name: String, key_suffix: String, scope: String, created_at: DateTime, last_used_at: Option>, expires_at: Option>, } impl From for ApiKeyInfo { fn from(key: ApiKeyProperties) -> Self { Self { uid: key.uid.into_inner(), name: key.name, key_suffix: key.key_suffix, scope: key.owner_type.to_string(), created_at: key.created_at.utc(), last_used_at: key.last_used_at.map(|t| t.utc()), expires_at: key.expires_at.map(|t| t.utc()), } } } impl fmt::Display for ApiKeyInfo { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let name = &self.name; let uid = &self.uid; let created_at = self.created_at.format("%Y-%m-%d %H:%M:%S UTC"); write!(f, "{name} ({uid}, created {created_at})") } } impl TableFormat for ApiKeyInfo { fn header() -> Vec { vec![ Cell::new("UID"), Cell::new("Name"), Cell::new("Key"), Cell::new("Scope"), Cell::new("Created"), Cell::new("Last Used"), Cell::new("Expires At"), ] } fn row(&self) -> Vec { vec![ Cell::new(&self.uid), Cell::new(&self.name), Cell::new(format!("wk-**{}", self.key_suffix)), Cell::new(&self.scope), Cell::new(format_approx_duration_from_now_utc(self.created_at)), Cell::new( self.last_used_at .map(format_approx_duration_from_now_utc) .unwrap_or_else(|| "Never".to_string()), ), Cell::new( self.expires_at .map(|dt| dt.format("%Y-%m-%d %H:%M:%S UTC").to_string()) .unwrap_or_else(|| "Never".to_string()), ), ] } } fn resolve_api_key_identifier( keys: &[ApiKeyInfo], key_identifier: &str, ) -> Result> { if let Some(key) = keys.iter().find(|key| key.uid == key_identifier) { return Ok(Some(key.clone())); } let mut matches = keys .iter() .filter(|key| key.name == key_identifier) .cloned() .collect::>(); matches.sort_by_key(|key| Reverse(key.created_at)); if matches.is_empty() { return Err(anyhow!("API key '{key_identifier}' not found")); } else if matches.len() == 1 { return Ok(Some(matches[0].clone())); } if io::stdin().is_terminal() { return match Select::new( &format!("Multiple API keys match '{key_identifier}'. Select a key to expire:"), matches, ) .prompt() { Ok(key) => Ok(Some(key)), Err(InquireError::OperationCanceled | InquireError::OperationInterrupted) => Ok(None), Err(err) => Err(err.into()), }; } println!("Multiple API keys match '{key_identifier}':"); for key in matches { println!(" {key}"); } Err(anyhow!( "Multiple API keys match '{key_identifier}'; specify the key by UID" )) } #[derive(Debug, Clone, Serialize)] struct CreatedApiKeyInfo { raw_api_key: String, api_key: ApiKeyInfo, } #[derive(Debug, Clone, Serialize)] struct ExpiredApiKeyInfo { key_uid: String, expired: bool, } fn sort_api_keys( keys: &mut [ApiKeyInfo], sort_by: Option, sort_order: Option, ) { let Some(sort_by) = sort_by else { return; }; let descending = matches!(sort_order, Some(SortOrderArg::Desc)); match sort_by { ApiKeySortByArg::Name => { if descending { keys.sort_by_key(|k| Reverse(k.name.to_lowercase())); } else { keys.sort_by_key(|k| k.name.to_lowercase()); } } ApiKeySortByArg::CreatedAt => { if descending { keys.sort_by_key(|k| Reverse(k.created_at)); } else { keys.sort_by_key(|k| k.created_at); } } ApiKeySortByArg::LastUsedAt => { if descending { keys.sort_by_key(|k| Reverse(k.last_used_at)); } else { keys.sort_by_key(|k| k.last_used_at); } } ApiKeySortByArg::ExpiresAt => { if descending { keys.sort_by_key(|k| Reverse(k.expires_at)); } else { keys.sort_by_key(|k| k.expires_at); } } ApiKeySortByArg::Scope => { if descending { keys.sort_by_key(|k| Reverse(k.scope.clone())); } else { keys.sort_by_key(|k| k.scope.clone()); } } } } fn expires_at_from_args(args: ApiKeyExpirationArgs) -> Result> { if args.no_expiration { return Ok(None); } if let Some(expires_at) = args.expires_at { return Ok(Some(Time::from(expires_at))); } if let Some(expires_in) = args.expires_in { let duration = chrono::Duration::from_std(expires_in.into()) .map_err(|_| anyhow!("expiration duration is too large"))?; return Ok(Some(Time::from(Utc::now() + duration))); } Err(anyhow!("expiration behavior is required")) } fn print_created_api_key( result: CreatedApiKeyInfo, output_format: OutputFormat, json_output: warp_cli::json_filter::JsonOutput, ) -> Result<()> { if json_output.force_json_output() { output::print_raw_json(serde_json::to_value(&result)?, &json_output)?; return Ok(()); } match output_format { OutputFormat::Json => output::write_json(&result, std::io::stdout())?, OutputFormat::Ndjson => output::write_json_line(&result, std::io::stdout())?, OutputFormat::Pretty | OutputFormat::Text => { println!("API key '{}' created.", result.api_key.name); println!("UID: {}", result.api_key.uid); println!("Raw API key: {}", result.raw_api_key); println!("This secret key is shown only once. Store it securely."); } } Ok(()) } fn print_expire_api_key_result( key_uid: String, expired: bool, output_format: OutputFormat, json_output: warp_cli::json_filter::JsonOutput, ) -> Result<()> { let result = ExpiredApiKeyInfo { key_uid, expired }; if json_output.force_json_output() { output::print_raw_json(serde_json::to_value(&result)?, &json_output)?; return Ok(()); } match output_format { OutputFormat::Json => output::write_json(&result, std::io::stdout())?, OutputFormat::Ndjson => output::write_json_line(&result, std::io::stdout())?, OutputFormat::Pretty | OutputFormat::Text => { if expired { println!("API key '{}' expired.", result.key_uid); } else { println!("API key '{}' was not expired.", result.key_uid); } } } Ok(()) } fn finish_command(result: Result<()>, ctx: &mut ModelContext) { match result { Ok(()) => ctx.terminate_app(TerminationMode::ForceTerminate, None), Err(err) => super::report_fatal_error(err, ctx), } } #[cfg(test)] #[path = "api_key_tests.rs"] mod tests;