creddy/src-tauri/creddy_cli/src/cli.rs

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<Action>,
}

impl Cli {
    // proxy the Parser method so that main crate doesn't have to depend on Clap
    pub fn parse() -> Self {
        <Self as Parser>::parse()
    }
}


#[derive(Debug, Clone, Args)]
pub struct GlobalArgs {
    /// Connect to the main Creddy application at this path
    #[arg(long, short = 'a')]
    server_addr: Option<PathBuf>,
}


#[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<String>,
    /// 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<String>,
}


#[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<PathBuf>,
    req: &CliRequest
) -> anyhow::Result<Result<CliResponse, ServerError>> {
    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<CliResponse, ServerError> = serde_json::from_slice(&buf)?;
    Ok(res)
}