creddy/src-tauri/src/cli.rs

use std::ffi::OsString;
use std::path::PathBuf;
use std::process::Command as ChildCommand;
#[cfg(windows)]
use std::time::Duration;

use clap::{
    Command,
    Arg,
    ArgMatches,
    ArgAction,
    builder::PossibleValuesParser,
    value_parser,
 };
use tokio::io::{AsyncReadExt, AsyncWriteExt};

use crate::errors::*;
use crate::srv::{
    self,
    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")
        .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) -> Result<(), CliError> {
    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, &Request::GetAwsCredentials { name, base })? {
        Response::AwsBase(creds) => serde_json::to_string(&creds).unwrap(),
        Response::AwsSession(creds) => serde_json::to_string(&creds).unwrap(),
        r => return Err(RequestError::Unexpected(r).into()),
    };
    println!("{output}");
    Ok(())
}


pub fn exec(args: &ArgMatches, global_args: &ArgMatches) -> Result<(), CliError> {
    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 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 make_request(addr, &Request::GetAwsCredentials { name, base })? {
        Response::AwsBase(creds) => {
            cmd.env("AWS_ACCESS_KEY_ID", creds.access_key_id);
            cmd.env("AWS_SECRET_ACCESS_KEY", creds.secret_access_key);
        },
        Response::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 => return Err(RequestError::Unexpected(r).into()),
    }

    #[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, global_args: &ArgMatches) -> Result<(), CliError> {
    let addr = global_args.get_one("server_addr").cloned();
    let action = match args.get_one::<String>("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(addr, &req) {
        Ok(Response::Empty) => Ok(()),
        Ok(r) => Err(RequestError::Unexpected(r).into()),
        Err(e) => Err(e.into()),
    }
}


#[tokio::main]
async fn make_request(addr: Option<PathBuf>, req: &Request) -> Result<Response, RequestError> {
    let mut data = serde_json::to_string(req).unwrap();
    // server expects newline marking end of request
    data.push('\n');

    let mut stream = 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<Response, ServerError> = serde_json::from_slice(&buf)?;
    Ok(res?)
}


#[cfg(windows)]
async fn connect(addr: Option<PathBuf>) -> Result<NamedPipeClient, std::io::Error> {
    // apparently attempting to connect can fail if there's already a client connected
    loop {
        let addr = addr.unwrap_or_else(|| srv::addr("creddy-server"));
        match ClientOptions::new().open(&addr) {
            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(addr: Option<PathBuf>) -> Result<UnixStream, std::io::Error> {
    let path = addr.unwrap_or_else(|| srv::addr("creddy-server"));
    UnixStream::connect(&path).await
}