creddy/src-tauri/src/cli.rs

use std::ffi::OsString;
use std::process::Command as ChildCommand;
#[cfg(unix)]
use std::os::unix::process::CommandExt;

use clap::{
    Command,
     Arg,
     ArgMatches,
     ArgAction
 };
use tokio::{
    net::TcpStream,
    io::{AsyncReadExt, AsyncWriteExt},
};


use crate::app;
use crate::config::AppConfig;
use crate::credentials::{BaseCredentials, SessionCredentials};
use crate::errors::*;


pub fn parser() -> Command<'static> {
    Command::new("creddy")
        .about("A friendly AWS credentials manager")
        .subcommand(
            Command::new("run")
                .about("Launch Creddy")
        )
        .subcommand(
            Command::new("show")
                .about("Fetch and display AWS credentials")
                .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)
                )
        )
}


pub fn show(args: &ArgMatches) -> Result<(), CliError> {
    let base = args.get_one("base").unwrap_or(&false);
    let creds = get_credentials(*base)?;
    println!("{creds}");
    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);
    
    if base {
        let creds: BaseCredentials = serde_json::from_str(&get_credentials(base)?)
            .map_err(|_| RequestError::InvalidJson)?;
        cmd.env("AWS_ACCESS_KEY_ID", creds.access_key_id);
        cmd.env("AWS_SECRET_ACCESS_KEY", creds.secret_access_key);
    }
    else {
        let creds: SessionCredentials = serde_json::from_str(&get_credentials(base)?)
            .map_err(|_| RequestError::InvalidJson)?;
        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.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));
    };
}


#[tokio::main]
async fn get_credentials(base: bool) -> Result<String, RequestError> {
    let pool = app::connect_db().await?;
    let config = AppConfig::load(&pool).await?;
    let path = if base {"/creddy/base-credentials"} else {"/"};

    let mut stream = TcpStream::connect((config.listen_addr, config.listen_port)).await?;
    let req = format!("GET {path} HTTP/1.0\r\n\r\n");
    stream.write_all(req.as_bytes()).await?;

    // some day we'll have a proper HTTP parser
    let mut buf = vec![0; 8192];
    stream.read_to_end(&mut buf).await?;

    let status = buf.split(|&c| &[c] == b" ")
        .skip(1)
        .next()
        .ok_or(RequestError::MalformedHttpResponse)?;

    if status != b"200" {
        let s = String::from_utf8_lossy(status).to_string();
        return Err(RequestError::Failed(s));
    }

    let break_idx = buf.windows(4)
        .position(|w| w == b"\r\n\r\n")
        .ok_or(RequestError::MalformedHttpResponse)?;
    let body = &buf[(break_idx + 4)..];

    let creds_str = std::str::from_utf8(body)
        .map_err(|_| RequestError::MalformedHttpResponse)?
        .to_string();

    if creds_str == "Denied!" {
        return Err(RequestError::Rejected);
    }
    Ok(creds_str)
}