use clap::{Parser, Subcommand};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::env;
use std::fs;
use std::io::{self, Write};
use std::path::PathBuf;
use std::process::Command;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
mod aws;
use aws::{AwsManager, ProfileInfo};
const BOLD_WHITE: &str = "\x1b[1;37m";
const RESET: &str = "\x1b[0m";
const KEE_ART: &str = r#"
██╗ ██╗███████╗███████╗
██║ ██╔╝██╔════╝██╔════╝
█████╔╝ █████╗ █████╗
██╔═██╗ ██╔══╝ ██╔══╝
██║ ██╗███████╗███████╗
╚═╝ ╚═╝╚══════╝╚══════╝
AWS CLI profile manager"#;
// Environment variable names
const KEE_ACTIVE_PROFILE: &str = "KEE_ACTIVE_PROFILE";
const KEE_CURRENT_PROFILE: &str = "KEE_CURRENT_PROFILE";
const AWS_PROFILE: &str = "AWS_PROFILE";
const AWS_CLI_AUTO_PROMPT: &str = "AWS_CLI_AUTO_PROMPT";
const AWS_PAGER: &str = "AWS_PAGER";
#[derive(Parser)]
#[command(name = "kee")]
#[command(version = env!("CARGO_PKG_VERSION"))]
#[command(about = format!("{KEE_ART}\n V. {}", env!("CARGO_PKG_VERSION")))]
#[command(long_about = format!("{KEE_ART}\n V. {}\n\nExamples:\n kee Show current profile or this help\n kee add myprofile Add a new AWS profile\n kee use Pick a profile interactively (starts sub-shell)\n kee use myprofile Use a specific profile (starts sub-shell)\n kee aws myprofile s3 ls Run an AWS CLI command with a profile\n kee run myprofile -- CMD Run any command with a profile (no sub-shell)\n kee ls List all available profiles\n kee current Show current, active profile\n kee rm Pick a profile to remove interactively\n kee rm myprofile Remove a specific profile", env!("CARGO_PKG_VERSION")))]
struct Cli {
#[command(subcommand)]
command: Option<Commands>,
}
#[derive(Subcommand)]
enum Commands {
/// Add a new AWS profile
Add {
#[arg(value_name = "PROFILE_NAME", help = "Name for the new AWS profile")]
profile_name: String,
},
/// Use an available profile
Use {
#[arg(
value_name = "PROFILE_NAME",
help = "Name of the AWS profile to use (interactive picker if omitted)"
)]
profile_name: Option<String>,
},
/// List all available profiles
Ls {
/// Only show profile names (useful for scripting)
#[arg(long)]
names: bool,
},
/// Show current active profile
Current,
/// Remove a profile
Rm {
#[arg(
value_name = "PROFILE_NAME",
help = "Name of the AWS profile to remove (interactive picker if omitted)"
)]
profile_name: Option<String>,
},
/// Run a command with a profile's credentials (no sub-shell)
Run {
#[arg(
value_name = "PROFILE_NAME",
help = "Name of the AWS profile to use for the command"
)]
profile_name: String,
#[arg(
value_name = "COMMAND",
trailing_var_arg = true,
allow_hyphen_values = true,
help = "Command and arguments to run (use -- before flags)"
)]
cmd: Vec<String>,
},
/// Run an AWS CLI command with a profile's credentials (sugar over `run`)
Aws {
#[arg(
value_name = "PROFILE_NAME",
help = "Name of the AWS profile to use for the command"
)]
profile_name: String,
#[arg(
value_name = "ARGS",
trailing_var_arg = true,
allow_hyphen_values = true,
help = "Arguments to pass to 'aws' (e.g. s3 ls)"
)]
args: Vec<String>,
},
/// Update profile settings
Set {
#[arg(
value_name = "PROFILE_NAME",
help = "Name of the AWS profile to update"
)]
profile_name: String,
/// Mark as a production account
#[arg(long, conflicts_with = "no_production")]
production: bool,
/// Unmark as a production account
#[arg(long)]
no_production: bool,
},
}
#[derive(Serialize, Deserialize, Default)]
struct KeeConfig {
profiles: HashMap<String, ProfileInfo>,
current_profile: Option<String>,
}
struct KeeManager {
config_file: PathBuf,
aws_manager: AwsManager,
}
fn hlt(text: &str) -> String {
format!("{BOLD_WHITE}{text}{RESET}")
}
const SPINNER_FRAMES: &[&str] = &["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
struct Spinner {
running: Arc<AtomicBool>,
handle: Option<thread::JoinHandle<()>>,
}
impl Spinner {
fn start(message: &str) -> Self {
let running = Arc::new(AtomicBool::new(true));
let running_clone = Arc::clone(&running);
let message = message.to_string();
let handle = thread::spawn(move || {
let mut i = 0;
while running_clone.load(Ordering::Relaxed) {
print!(
"\r {} {}",
SPINNER_FRAMES[i % SPINNER_FRAMES.len()],
message
);
let _ = io::stdout().flush();
thread::sleep(Duration::from_millis(80));
i += 1;
}
});
Self {
running,
handle: Some(handle),
}
}
fn stop(mut self, result: &str) {
self.running.store(false, Ordering::Relaxed);
if let Some(handle) = self.handle.take() {
let _ = handle.join();
}
// Clear the spinner line and print the result
print!("\r\x1b[2K");
let _ = io::stdout().flush();
println!("{result}");
}
}
impl Drop for Spinner {
fn drop(&mut self) {
self.running.store(false, Ordering::Relaxed);
if let Some(handle) = self.handle.take() {
let _ = handle.join();
}
print!("\r\x1b[2K");
let _ = io::stdout().flush();
}
}
/// Buffer before token expiry at which we refresh, in seconds.
const REFRESH_BUFFER_SECS: i64 = 300;
/// Fallback sleep duration if expiry can't be read, in seconds.
const REFRESH_FALLBACK_SECS: u64 = 1800;
/// Background thread that keeps the SSO session fresh while a sub-shell is active.
struct SessionRefresher {
running: Arc<AtomicBool>,
handle: Option<thread::JoinHandle<()>>,
}
impl SessionRefresher {
fn start(aws_manager: AwsManager, profile_info: ProfileInfo) -> Self {
let running = Arc::new(AtomicBool::new(true));
let running_clone = Arc::clone(&running);
let handle = thread::spawn(move || {
let mut last_ok = true;
while running_clone.load(Ordering::Relaxed) {
// Decide how long to sleep based on the current token expiry.
let sleep_secs = match aws_manager.read_token_expiry(&profile_info) {
Some(expires_at) => {
let remaining =
(expires_at - chrono::Utc::now()).num_seconds() - REFRESH_BUFFER_SECS;
// If we're already within the buffer, refresh immediately.
remaining.max(0) as u64
}
None => REFRESH_FALLBACK_SECS,
};
// Sleep in short ticks so we exit promptly when signalled.
let mut slept = 0u64;
while slept < sleep_secs && running_clone.load(Ordering::Relaxed) {
thread::sleep(Duration::from_secs(1));
slept += 1;
}
if !running_clone.load(Ordering::Relaxed) {
break;
}
// Refresh. Surface state transitions to the user via stderr so a
// silently-dying session is at least visible in the sub-shell.
let ok = aws_manager.try_refresh_token(&profile_info);
match (last_ok, ok) {
(true, false) => {
eprintln!(
"\n Kee: background session refresh failed for '{}'.\n Run 'aws sso login --profile {}' if AWS calls start failing.",
profile_info.profile_name, profile_info.profile_name
);
}
(false, true) => {
eprintln!(
"\n Kee: background session refresh recovered for '{}'.",
profile_info.profile_name
);
}
_ => {}
}
last_ok = ok;
}
});
Self {
running,
handle: Some(handle),
}
}
}
impl Drop for SessionRefresher {
fn drop(&mut self) {
self.running.store(false, Ordering::Relaxed);
if let Some(handle) = self.handle.take() {
let _ = handle.join();
}
}
}
impl KeeManager {
fn prompt_user(&self, message: &str) -> io::Result<bool> {
print!("{message}");
io::stdout().flush()?;
let mut input = String::new();
io::stdin().read_line(&mut input)?;
Ok(input.trim().to_lowercase() == "y")
}
/// Show a fuzzy picker over configured profiles. Returns the chosen profile
/// name, or None if there are no profiles or the user aborted.
fn pick_profile(&self, prompt: &str) -> io::Result<Option<String>> {
let config = self.load_config();
if config.profiles.is_empty() {
println!(
"\n [!] No profiles configured.\n Run {} to add one.",
hlt("kee add PROFILE_NAME")
);
return Ok(None);
}
// Sorted for stable ordering across invocations.
let mut names: Vec<String> = config.profiles.keys().cloned().collect();
names.sort();
let selection =
dialoguer::FuzzySelect::with_theme(&dialoguer::theme::ColorfulTheme::default())
.with_prompt(prompt)
.items(&names)
.default(0)
.interact_opt()
.map_err(io::Error::other)?;
Ok(selection.map(|i| names[i].clone()))
}
fn new() -> io::Result<Self> {
let home_dir = dirs::home_dir().ok_or_else(|| {
io::Error::new(
io::ErrorKind::NotFound,
"\n [X] Could not find home directory\n",
)
})?;
let config_dir = home_dir.join(".kee");
let config_file = config_dir.join("config.json");
// Create .kee directory if it doesn't exist
fs::create_dir_all(&config_dir)?;
let aws_manager = AwsManager::new()?;
Ok(Self {
config_file,
aws_manager,
})
}
fn load_config(&self) -> KeeConfig {
if !self.config_file.exists() {
return KeeConfig::default();
}
match fs::read_to_string(&self.config_file) {
Ok(content) => serde_json::from_str(&content).unwrap_or_default(),
Err(_) => KeeConfig::default(),
}
}
fn save_config(&self, config: &KeeConfig) -> io::Result<()> {
let content = serde_json::to_string_pretty(config)?;
fs::write(&self.config_file, content)
}
fn add_profile(&self, profile_name: &str) -> io::Result<bool> {
println!("\n Starting SSO configuration...");
println!(" (This will open your browser to complete authentication.)");
println!("\n Follow the prompts:");
println!(" {} Enter your SSO start URL", hlt("1."));
println!(" {} Enter your SSO region", hlt("2."));
println!(" {} Authenticate in your browser", hlt("3."));
println!(" {} Select your AWS account", hlt("4."));
println!(" {} Select your role", hlt("5."));
println!(
" {} Choose your output format (recommend: json)",
hlt("6.")
);
println!(
"\n {} A session can be liked to multiple profiles.\n When prompted for a 'session name', use something generic, like your company name.\n",
hlt("Tip:")
);
// Run aws configure sso
let status = Command::new("aws")
.args(["configure", "sso", "--profile", profile_name])
.status()?;
if !status.success() {
println!(" [X] SSO configuration failed.");
return Ok(false);
}
println!(
"\n {} You can ignore the AWS CLI example above.\n {} will handle profiles for you.",
hlt("Note:"),
hlt("Kee")
);
// Reformat the AWS config file
self.aws_manager.format_config()?;
// Read profile info
let profile_info = match self.aws_manager.read_profile(profile_name) {
Some(info) => info,
None => {
println!("\n [X] Could not read profile information.");
return Ok(false);
}
};
// Save to kee config
let mut config = self.load_config();
// Ask if this is a production account
let production = self.prompt_user(&format!(
"\n Is {} a production account? (y/N): ",
hlt(profile_name)
))?;
let mut profile_info = profile_info;
profile_info.production = production;
config
.profiles
.insert(profile_name.to_string(), profile_info);
self.save_config(&config)?;
// Test the profile
if self.check_credentials(profile_name) {
println!("\n [✓] The profile was added and it's working!");
} else {
println!("\n [X] I created the profile but credentials may need a refresh...");
println!(" {} aws sso login --profile {}", hlt("Try:"), profile_name);
}
Ok(true)
}
fn list_profiles(&self, names: bool) {
let config = self.load_config();
if config.profiles.is_empty() {
if !names {
println!(
"\n [!] No profiles configured.\n Run {} to add one.",
hlt("kee add PROFILE_NAME")
);
}
return;
}
if names {
for profile_name in config.profiles.keys() {
println!("{profile_name}");
}
return;
}
println!();
for (profile_name, profile_info) in &config.profiles {
let status = if Some(profile_name.as_str()) == config.current_profile.as_deref() {
" (Current profile)"
} else {
""
};
println!(" {}{}", hlt(profile_name), status);
println!(" • {} {}", hlt("Account ID:"), profile_info.sso_account_id);
println!(" • {} {}\n", hlt("Role:"), profile_info.sso_role_name);
}
}
fn remove_profile(&self, profile_name: &str) -> io::Result<bool> {
let mut config = self.load_config();
if !config.profiles.contains_key(profile_name) {
println!("\n [!] Profile '{}' not found.", hlt(profile_name));
return Ok(false);
}
// Confirm removal
if !self.prompt_user(&format!(
"\n [!] Are you sure you want to remove profile '{}'? (y/N): ",
hlt(profile_name)
))? {
return Ok(false);
}
// Get profile info before removal
let profile_info = config.profiles.remove(profile_name).unwrap();
// Clear current profile if it's the one being removed
if config.current_profile.as_deref() == Some(profile_name) {
config.current_profile = None;
}
self.save_config(&config)?;
// Remove the AWS profile from config file
let hlt_profile = hlt(profile_name);
match self.aws_manager.remove_profile(&profile_info.profile_name) {
Ok(_) => {
println!(" [✓] Profile '{hlt_profile}' has been removed.");
}
Err(e) => {
println!(" [✓] Profile '{hlt_profile}' removed from {}.", hlt("Kee"));
println!(
" [!] Could not remove AWS profile '{}': {}",
hlt(&profile_info.profile_name),
e
);
println!(
" You may want to remove it manually from {}",
hlt("~/.aws/config")
);
}
}
Ok(true)
}
fn set_profile(
&self,
profile_name: &str,
production: bool,
no_production: bool,
) -> io::Result<bool> {
let mut config = self.load_config();
let profile = match config.profiles.get_mut(profile_name) {
Some(p) => p,
None => {
println!("\n [!] Profile '{}' not found.", hlt(profile_name));
return Ok(false);
}
};
if production {
profile.production = true;
println!(
"\n [✓] Profile '{}' marked as production.",
hlt(profile_name)
);
} else if no_production {
profile.production = false;
println!(
"\n [✓] Profile '{}' unmarked as production.",
hlt(profile_name)
);
}
self.save_config(&config)?;
Ok(true)
}
fn use_profile(&self, profile_name: &str) -> io::Result<bool> {
// Check if already in a Kee profile
if env::var(KEE_ACTIVE_PROFILE).is_ok() {
let current_profile =
env::var(KEE_CURRENT_PROFILE).unwrap_or_else(|_| "unknown".to_string());
println!(
"\n [!] You are using a {} profile: {}",
hlt("Kee"),
hlt(¤t_profile)
);
println!(" Exit the current session first by typing {}", hlt("exit"));
return Ok(false);
}
let mut config = self.load_config();
let hlt_profile = hlt(profile_name);
if !config.profiles.contains_key(profile_name) {
println!("\n [!] Profile '{hlt_profile}' not found.");
if !config.profiles.is_empty() {
println!(" Available profiles:");
for name in config.profiles.keys() {
println!(" • {}\n", hlt(name));
}
}
// Offer to add the profile
if self.prompt_user(" Would you like to add now? (y/N): ")? {
if self.add_profile(profile_name)? {
if self.prompt_user(&format!(
" Would you like to use profile '{hlt_profile}' now? (y/N): "
))? {
// Reload config
config = self.load_config();
} else {
println!(
"\n Profile '{}' is ready to use. Run {} when needed.",
hlt_profile,
hlt(&format!("kee use {profile_name}"))
);
return Ok(true);
}
} else {
println!(" [X] Failed to add profile '{hlt_profile}'.");
return Ok(false);
}
} else {
return Ok(false);
}
}
let profile_info = config.profiles.get(profile_name).unwrap().clone();
let profile_name = profile_info.profile_name.clone();
// Always attempt a token refresh to maximise session duration
println!();
let spinner = Spinner::start("Refreshing session...");
let refreshed = self.aws_manager.try_refresh_token(&profile_info)
&& self.check_credentials(&profile_name);
if refreshed {
spinner.stop(" [✓] Session refreshed.");
} else if self.check_credentials(&profile_name) {
spinner.stop(" [✓] Session is valid.");
} else {
spinner.stop(" [!] Session expired. Opening SSO login...");
if !self.sso_login(&profile_name)? {
println!(
" [X] Failed to authenticate. Please run {} manually.",
hlt("aws sso login")
);
return Ok(false);
}
}
// Update current profile
config.current_profile = Some(profile_name.clone());
self.save_config(&config)?;
// Start subshell
self.start_subshell(&profile_info)?;
// Clear current profile when subshell exits
config.current_profile = None;
self.save_config(&config)?;
Ok(true)
}
/// Run a single command with the chosen profile's credentials and exit
/// with the command's exit code. No sub-shell. Output of the wrapped
/// command is left untouched (Kee's own status messages go to stderr).
fn run_command(&self, profile_name: &str, cmd: &[String]) -> io::Result<i32> {
if cmd.is_empty() {
eprintln!("\n [X] Please specify a command to run");
eprintln!(" Usage: {}", hlt("kee run PROFILE_NAME -- CMD ARGS..."));
return Ok(2);
}
let config = self.load_config();
let profile_info = match config.profiles.get(profile_name) {
Some(p) => p.clone(),
None => {
eprintln!("\n [!] Profile '{}' not found.", hlt(profile_name));
eprintln!(" Run {} to see available profiles.", hlt("kee ls"));
return Ok(1);
}
};
let aws_profile = profile_info.profile_name.clone();
// Ensure credentials are usable. Try a refresh, then validate, then
// fall back to interactive sso login. All status to stderr so we don't
// contaminate the wrapped command's stdout.
let _ = self.aws_manager.try_refresh_token(&profile_info);
if !self.check_credentials(&aws_profile) {
eprintln!(
" Kee: session expired for '{}', running 'aws sso login'...",
hlt(&aws_profile)
);
if !self.sso_login(&aws_profile)? || !self.check_credentials(&aws_profile) {
eprintln!(
" [X] Failed to authenticate. Please run {} manually.",
hlt("aws sso login")
);
return Ok(1);
}
}
if profile_info.production {
eprintln!(" \x1b[1;31m⚠️ PRODUCTION ACCOUNT ({})\x1b[0m", aws_profile);
}
let (program, args) = cmd.split_first().unwrap();
let status = Command::new(program)
.args(args)
.env(AWS_PROFILE, &aws_profile)
.env(KEE_CURRENT_PROFILE, &aws_profile)
.env(KEE_ACTIVE_PROFILE, "1")
.status();
match status {
Ok(s) => Ok(s.code().unwrap_or(1)),
Err(e) => {
eprintln!(" [X] Failed to execute '{}': {}", program, e);
Ok(127)
}
}
}
fn current_profile(&self) {
// Check if in active session
if let Ok(current) = env::var(KEE_CURRENT_PROFILE) {
println!("\n Current profile: {}", hlt(¤t));
println!(" Type {} to return to your main shell.", hlt("exit"));
} else {
let config = self.load_config();
match config.current_profile {
Some(current) => println!("\n Current profile: {}", hlt(¤t)),
None => println!("\n [!] No profile is currently active."),
}
}
}
fn check_credentials(&self, profile_name: &str) -> bool {
match Command::new("aws")
.args(["sts", "get-caller-identity", "--profile", profile_name])
.env(AWS_CLI_AUTO_PROMPT, "off")
.env(AWS_PAGER, "")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status()
{
Ok(status) => status.success(),
Err(_) => false,
}
}
fn sso_login(&self, profile_name: &str) -> io::Result<bool> {
let status = Command::new("aws")
.args(["sso", "login", "--profile", profile_name])
.status()?;
Ok(status.success())
}
fn start_subshell(&self, profile_info: &ProfileInfo) -> io::Result<()> {
let profile_name = &profile_info.profile_name;
// Get current shell
let shell = if cfg!(windows) {
env::var("COMSPEC").unwrap_or_else(|_| "cmd.exe".to_string())
} else {
env::var("SHELL").unwrap_or_else(|_| "/bin/bash".to_string())
};
// Show banner
if profile_info.production {
println!("\n \x1b[1;31m⚠️ PRODUCTION ACCOUNT\x1b[0m");
}
println!("\n Profile: {}", hlt(profile_name));
println!(" {} is starting a sub-shell...", hlt("Kee"));
println!(" Type {} to return to your main shell.", hlt("exit"));
// Start subshell with environment
let mut cmd = Command::new(&shell);
cmd.env(AWS_PROFILE, profile_name);
cmd.env(KEE_CURRENT_PROFILE, profile_name);
cmd.env(KEE_ACTIVE_PROFILE, "1");
// Update PS1 for Unix-like systems
if !cfg!(windows) {
if let Ok(ps1) = env::var("PS1") {
cmd.env("PS1", format!("aws:{profile_name} {ps1}"));
} else {
cmd.env("PS1", format!("aws:{profile_name} $ "));
}
}
// Keep the session refreshed in the background while the sub-shell is alive.
// Dropped automatically when this function returns.
let _refresher = SessionRefresher::start(self.aws_manager.clone(), profile_info.clone());
let _ = cmd.status();
println!("\n {} — Session ended.", hlt(profile_name));
Ok(())
}
}
fn main() -> io::Result<()> {
let cli = match Cli::try_parse() {
Ok(cli) => cli,
Err(err) => {
// Customise the missing-arg error for `kee add` only; `use` and `rm`
// accept an optional name and fall through to a picker.
if err.kind() == clap::error::ErrorKind::MissingRequiredArgument {
let error_msg = err.to_string();
if error_msg.contains("<PROFILE_NAME>") && error_msg.contains("kee add") {
eprintln!("\n [X] Please specify a name for the new profile");
eprintln!(" Usage: {}", hlt("kee add PROFILE_NAME"));
std::process::exit(2);
}
}
// For all other errors, use clap's default handling
err.exit();
}
};
let kee = KeeManager::new()?;
// No subcommand: show the current session if active, otherwise print help.
// Only KEE_ACTIVE_PROFILE counts as "in a session" — the config field can
// go stale if a sub-shell exits abnormally.
let command = match cli.command {
Some(c) => c,
None => {
if env::var(KEE_ACTIVE_PROFILE).is_ok() {
kee.current_profile();
return Ok(());
}
// Re-parse with --help so clap prints and exits.
let _ = Cli::try_parse_from(["kee", "--help"])
.err()
.map(|e| e.exit());
return Ok(());
}
};
match command {
Commands::Add { profile_name } => {
kee.add_profile(&profile_name)?;
}
Commands::Use { profile_name } => {
let name = match profile_name {
Some(n) => n,
None => match kee.pick_profile("Which profile?")? {
Some(n) => n,
None => return Ok(()),
},
};
kee.use_profile(&name)?;
}
Commands::Ls { names } => {
kee.list_profiles(names);
}
Commands::Current => {
kee.current_profile();
}
Commands::Rm { profile_name } => {
let name = match profile_name {
Some(n) => n,
None => match kee.pick_profile("Which profile?")? {
Some(n) => n,
None => return Ok(()),
},
};
kee.remove_profile(&name)?;
}
Commands::Set {
profile_name,
production,
no_production,
} => {
kee.set_profile(&profile_name, production, no_production)?;
}
Commands::Run { profile_name, cmd } => {
let code = kee.run_command(&profile_name, &cmd)?;
std::process::exit(code);
}
Commands::Aws { profile_name, args } => {
// Sugar: prepend "aws" and dispatch through run_command.
let mut cmd = Vec::with_capacity(args.len() + 1);
cmd.push("aws".to_string());
cmd.extend(args);
let code = kee.run_command(&profile_name, &cmd)?;
std::process::exit(code);
}
}
Ok(())
}