use std::env; 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::{ Command, Arg, ArgMatches, ArgAction, builder::PossibleValuesParser, value_parser, }; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::proto::{ CliCredential, CliRequest, CliResponse, ServerError, ShortcutAction, }; pub fn parser() -> Command<'static> { Command::new("creddy") .version(env!("CARGO_PKG_VERSION")) .about("A friendly credential manager") .arg( Arg::new("server_addr") .short('a') .long("server-addr") .takes_value(true) .value_parser(value_parser!(PathBuf)) .help("Connect to the main Creddy process at this address") ) .subcommand( Command::new("run") .about("Launch Creddy") ) .subcommand( Command::new("get") .about("Request AWS credentials from Creddy and output to stdout") .arg( Arg::new("base") .short('b') .long("base") .action(ArgAction::SetTrue) .help("Use base credentials instead of session credentials") ) .arg( Arg::new("name") .help("If unspecified, use default credentials") ) ) .subcommand( Command::new("exec") .about("Inject AWS credentials into the environment of another command") .trailing_var_arg(true) .arg( Arg::new("base") .short('b') .long("base") .action(ArgAction::SetTrue) .help("Use base credentials instead of session credentials") ) .arg( Arg::new("name") .short('n') .long("name") .takes_value(true) .help("If unspecified, use default credentials") ) .arg( Arg::new("command") .multiple_values(true) ) ) .subcommand( Command::new("shortcut") .about("Invoke an action normally trigged by hotkey (e.g. launch terminal)") .arg( Arg::new("action") .value_parser( PossibleValuesParser::new(["show_window", "launch_terminal"]) ) ) ) } pub fn get(args: &ArgMatches, global_args: &ArgMatches) -> anyhow::Result<()> { let name = args.get_one("name").cloned(); let base = *args.get_one("base").unwrap_or(&false); let addr = global_args.get_one("server_addr").cloned(); let output = match make_request(addr, &CliRequest::GetCredential { name, base })?? { CliResponse::Credential(c) => serde_json::to_string_pretty(&c).unwrap(), r => bail!("Unexpected response from server: {r}"), }; println!("{output}"); Ok(()) } pub fn exec(args: &ArgMatches, global_args: &ArgMatches) -> anyhow::Result<()> { let name = args.get_one("name").cloned(); let base = *args.get_one("base").unwrap_or(&false); let addr = global_args.get_one("server_addr").cloned(); // Clap guarantees that cmd_line will be a sequence of at least 1 item // test this! let mut cmd_line = args.get_many("command").unwrap(); let cmd_name: &String = cmd_line.next().unwrap(); let mut cmd = ChildCommand::new(cmd_name); cmd.args(cmd_line); match make_request(addr, &CliRequest::GetCredential { name, base })?? { 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)] { // cmd.exec() never returns if successful let e = cmd.exec(); Err(e).with_context(|| { // eventually figure out how to display the actual command format!("Failed to execute command") })?; 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: &ArgMatches, global_args: &ArgMatches) -> anyhow::Result<()> { let addr = global_args.get_one("server_addr").cloned(); let action = match args.get_one::("action").map(|s| s.as_str()) { Some("show_window") => ShortcutAction::ShowWindow, Some("launch_terminal") => ShortcutAction::LaunchTerminal, Some(&_) | None => unreachable!("Unknown shortcut action"), // guaranteed by clap }; let req = CliRequest::InvokeShortcut(action); match make_request(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) }