use std::ffi::OsString; use std::process::Command as ChildCommand; #[cfg(windows)] use std::time::Duration; use clap::{ Command, Arg, ArgMatches, ArgAction, builder::PossibleValuesParser, }; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::credentials::Credentials; use crate::errors::*; use crate::server::{Request, Response}; use crate::shortcuts::ShortcutAction; #[cfg(unix)] use { std::os::unix::process::CommandExt, tokio::net::UnixStream, }; #[cfg(windows)] use { tokio::net::windows::named_pipe::{NamedPipeClient, ClientOptions}, windows::Win32::Foundation::ERROR_PIPE_BUSY, }; pub fn parser() -> Command<'static> { Command::new("creddy") .version(env!("CARGO_PKG_VERSION")) .about("A friendly AWS credentials manager") .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") ) ) .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("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) -> Result<(), CliError> { let base = args.get_one("base").unwrap_or(&false); let output = match get_credentials(*base)? { Credentials::Base(creds) => serde_json::to_string(&creds).unwrap(), Credentials::Session(creds) => serde_json::to_string(&creds).unwrap(), }; println!("{output}"); Ok(()) } pub fn exec(args: &ArgMatches) -> Result<(), CliError> { let base = *args.get_one("base").unwrap_or(&false); let mut cmd_line = args.get_many("command") .ok_or(ExecError::NoCommand)?; let cmd_name: &String = cmd_line.next().unwrap(); // Clap guarantees that there will be at least one let mut cmd = ChildCommand::new(cmd_name); cmd.args(cmd_line); match get_credentials(base)? { Credentials::Base(creds) => { cmd.env("AWS_ACCESS_KEY_ID", creds.access_key_id); cmd.env("AWS_SECRET_ACCESS_KEY", creds.secret_access_key); }, Credentials::Session(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); } } #[cfg(unix)] { // cmd.exec() never returns if successful let e = cmd.exec(); match e.kind() { std::io::ErrorKind::NotFound => { let name: OsString = cmd_name.into(); Err(ExecError::NotFound(name).into()) } _ => Err(ExecError::ExecutionFailed(e).into()), } } #[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) -> Result<(), CliError> { 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 = Request::InvokeShortcut(action); match make_request(&req) { Ok(Response::Empty) => Ok(()), Ok(r) => Err(RequestError::Unexpected(r).into()), Err(e) => Err(e.into()), } } fn get_credentials(base: bool) -> Result { let req = Request::GetAwsCredentials { base }; match make_request(&req) { Ok(Response::Aws(creds)) => Ok(creds), Ok(r) => Err(RequestError::Unexpected(r)), Err(e) => Err(e), } } #[tokio::main] async fn make_request(req: &Request) -> Result { let mut data = serde_json::to_string(req).unwrap(); // server expects newline marking end of request data.push('\n'); let mut stream = connect().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?) } #[cfg(windows)] async fn connect() -> Result { // apparently attempting to connect can fail if there's already a client connected loop { match ClientOptions::new().open(r"\\.\pipe\creddy-requests") { Ok(stream) => return Ok(stream), Err(e) if e.raw_os_error() == Some(ERROR_PIPE_BUSY.0 as i32) => (), Err(e) => return Err(e), } tokio::time::sleep(Duration::from_millis(10)).await; } } #[cfg(unix)] async fn connect() -> Result { UnixStream::connect("/tmp/creddy.sock").await }