use std::path::PathBuf; use std::process::Command as ChildCommand; #[cfg(unix)] use std::os::unix::process::CommandExt; #[cfg(windows)] use std::time::Duration; use anyhow::{bail, Context}; use clap::{ Args, Parser, Subcommand }; use clap::builder::styling::{Styles, AnsiColor}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::proto::{ CliCredential, CliRequest, CliResponse, ServerError, ShortcutAction, }; #[derive(Debug, Parser)] #[command( about, version, name = "creddy", bin_name = "creddy", styles = Styles::styled() .header(AnsiColor::Yellow.on_default()) .usage(AnsiColor::Yellow.on_default()) .literal(AnsiColor::Green.on_default()) .placeholder(AnsiColor::Green.on_default()) )] /// A friendly credential manager pub struct Cli { #[command(flatten)] pub global_args: GlobalArgs, #[command(subcommand)] pub action: Option, } impl Cli { // proxy the Parser method so that main crate doesn't have to depend on Clap pub fn parse() -> Self { ::parse() } } #[derive(Debug, Clone, Args)] pub struct GlobalArgs { /// Connect to the main Creddy application at this path #[arg(long, short = 'a')] server_addr: Option, } #[derive(Debug, Subcommand)] pub enum Action { /// Launch Creddy Run, /// Request credentials from Creddy and output to stdout Get(GetArgs), /// Inject credentials into the environment of another command Exec(ExecArgs), /// Invoke an action normally triggered by hotkey (e.g. launch terminal) Shortcut(InvokeArgs), } #[derive(Debug, Args)] pub struct GetArgs { /// If unspecified, use default credentials #[arg(short, long)] name: Option, /// Use base credentials instead of session credentials (only applicable to AWS) #[arg(long, short, default_value_t = false)] base: bool, } #[derive(Debug, Args)] pub struct ExecArgs { #[command(flatten)] get_args: GetArgs, #[arg(trailing_var_arg = true)] /// Command to be wrapped command: Vec, } #[derive(Debug, Args)] pub struct InvokeArgs { #[arg(value_name = "ACTION", value_enum)] shortcut_action: ShortcutAction, } pub fn get(args: GetArgs, global: GlobalArgs) -> anyhow::Result<()> { let req = CliRequest::GetCredential { name: args.name, base: args.base, }; let output = match make_request(global.server_addr, &req)?? { CliResponse::Credential(CliCredential::AwsBase(c)) => { serde_json::to_string_pretty(&c).unwrap() }, CliResponse::Credential(CliCredential::AwsSession(c)) => { serde_json::to_string_pretty(&c).unwrap() }, r => bail!("Unexpected response from server: {r}"), }; println!("{output}"); Ok(()) } pub fn exec(args: ExecArgs, global: GlobalArgs) -> anyhow::Result<()> { // Clap guarantees that cmd_line will be a sequence of at least 1 item // test this! let mut cmd_line = args.command.iter(); let cmd_name = cmd_line.next().unwrap(); let mut cmd = ChildCommand::new(cmd_name); cmd.args(cmd_line); let req = CliRequest::GetCredential { name: args.get_args.name, base: args.get_args.base, }; match make_request(global.server_addr, &req)?? { CliResponse::Credential(CliCredential::AwsBase(creds)) => { cmd.env("AWS_ACCESS_KEY_ID", creds.access_key_id); cmd.env("AWS_SECRET_ACCESS_KEY", creds.secret_access_key); }, CliResponse::Credential(CliCredential::AwsSession(creds)) => { cmd.env("AWS_ACCESS_KEY_ID", creds.access_key_id); cmd.env("AWS_SECRET_ACCESS_KEY", creds.secret_access_key); cmd.env("AWS_SESSION_TOKEN", creds.session_token); }, r => bail!("Unexpected response from server: {r}"), } #[cfg(unix)] { let e = cmd.exec(); // cmd.exec() never returns if successful, so we never hit this line unless there's an error Err(e).with_context(|| { // eventually figure out how to display the actual command format!("Failed to execute command: {}", args.command.join(" ")) })?; Ok(()) } #[cfg(windows)] { let mut child = match cmd.spawn() { Ok(c) => c, Err(e) if e.kind() == std::io::ErrorKind::NotFound => { let name: OsString = cmd_name.into(); return Err(ExecError::NotFound(name).into()); } Err(e) => return Err(ExecError::ExecutionFailed(e).into()), }; let status = child.wait() .map_err(|e| ExecError::ExecutionFailed(e))?; std::process::exit(status.code().unwrap_or(1)); }; } pub fn invoke_shortcut(args: InvokeArgs, global: GlobalArgs) -> anyhow::Result<()> { let req = CliRequest::InvokeShortcut(args.shortcut_action); match make_request(global.server_addr, &req)?? { CliResponse::Empty => Ok(()), r => bail!("Unexpected response from server: {r}"), } } // Explanation for double-result: the server will return a (serialized) Result // to indicate when the operation succeeded or failed, which we deserialize. // However, the operation may fail to even communicate with the server, in // which case we return the outer Result #[tokio::main] async fn make_request( addr: Option, req: &CliRequest ) -> anyhow::Result> { let mut data = serde_json::to_string(req).unwrap(); // server expects newline marking end of request data.push('\n'); let mut stream = crate::connect(addr).await?; stream.write_all(&data.as_bytes()).await?; let mut buf = Vec::with_capacity(1024); stream.read_to_end(&mut buf).await?; let res: Result = serde_json::from_slice(&buf)?; Ok(res) }