index : bitwarden-ssh-agent.git

use std::{env, process::{Command, Stdio}, ffi::OsString};
use std::str::FromStr;
use gumdrop::Options;
use anyhow::{Result, Context};
use colored::*;
use crate::bwutil::BitwardenItem;

mod util;
mod bwutil;

/// Environment variable housing an existing Bitwarden session.
const SESSION_ENV_KEY: &str = "BW_SESSION";
const BW_SSH_DEBUG_ENV_KEY: &str = "BW_SSH_DEBUG";
/// Environment variable referencing the name of the field containing the SSH key passphrase (if any).
const BW_FIELD_KEY_PASSPHRASE: &str = "BW_KEY_PASSPHRASE";

/// A macro to print when the debug CLI arg is enabled.
macro_rules! debug_println {
    ($($arg:tt)*) => {
        if let Ok(enabled) = env::var(BW_SSH_DEBUG_ENV_KEY) {
            if bool::from_str(&enabled).unwrap_or(false) {
                println!("{}",format!($($arg)*).yellow());
            }
        }
    }
}

macro_rules! warn_println {
    ($($arg:tt)*) => {
        println!("{}",format!($($arg)*).yellow());
    }
}

macro_rules! info_println {
    ($($arg:tt)*) => {
        println!("{}",format!($($arg)*).green());
    }
}

fn main() -> Result<()> {
    // This environment variable only set when calling ssh-agent, so return the passphrase to authenticate and then quit.
    let unlock_passphrase = env::var(BW_FIELD_KEY_PASSPHRASE);
    if unlock_passphrase.is_ok() {
        println!("{}", unlock_passphrase.unwrap());
        return Ok(())
    }

    // Parse args and set the debug environment variable
    let args: util::Cli = util::Cli::parse_args_default_or_exit();

    env::set_var(BW_SSH_DEBUG_ENV_KEY, OsString::from(args.debug.to_string()));

    // Print the version information and exit if the --version flag was passed.
    if args.version {
        println!("{}", &util::get_version_string()?);
        return Ok(());
    }

    let session_token: String = check_session_token(&args)?;
    if !session_token.is_empty() {
        let folders = bwutil::exec_folder_search(&session_token, &args.folder)?;
        debug_println!("Found {} folder(s) named `{}`", folders.len().to_string().cyan(), args.folder.cyan());
        
        // Retrieve items from each folder since there may be multiple folders with the same name.
        // Pre-filter the items by checking for the presence of custom fields AND at least one attachment.
        for folder in folders {
            let folder_items: Vec<BitwardenItem> = bwutil::exec_list_folder_items(&session_token, &folder.id)?;
            let folder_items: Vec<&BitwardenItem> = folder_items
                .iter()
                .filter(|&item| item.fields.is_some() && item.attachments.is_some())
                .collect();
            
            debug_println!("Found {} item(s) in folder `{}` id({}) with at least one custom field and attachment", folder_items.len().to_string().cyan(), folder.name.cyan(), folder.id.cyan());
            
            for item in folder_items {
                // In order for this to be considered a valid SSH key item, this item needs to have the following fields:
                // - private - (required) this is the filename of the attachment containing the SSH key. 
                //		The field name is not case sensitive in order to be more user-friendly.
                // - passphrase - (optional) this is the passphrase (if any) for the SSH key. 
                //		The field name is not case sensitive in order to be more user-friendly.
                // - attachment (required) an attachment with the name specified in `private`.
                if let Some(fields) = &item.fields {
                    let key_filename: String = fields
                        .iter()
                        .find(|field| field.name.to_lowercase().eq(&args.key.to_lowercase()) && field.value.is_some())
                        .map(|field| field.value.as_ref().unwrap().clone())
                        .unwrap_or_default();

                    if !&key_filename.is_empty() {
                        let key_passphrase: String = fields
                            .iter()
                            .find(|field| field.name.to_lowercase().eq(&args.passphrase.to_lowercase()) && field.value.is_some())
                            .map(|field| field.value.as_ref().unwrap().clone())
                            .unwrap_or_default();

                        if let Some(attachments) = &item.attachments {
                            attachments
                            .iter()
                            .filter(|attachment| attachment.file_name.eq(&key_filename))
                            .for_each(|attachment| register_key(&item, &attachment.id, &key_passphrase, &session_token).unwrap());
                        }
                    }
                }
            }
        }
    }

    Ok(())
}

/// Gets a session token from the active Bitwarden session. The checking process is as follows:
/// 1. Check for `--session` argument and use that.
/// 2. Check for `BW_SESSION` environment variable and use that.
/// 3. Check for login status from `bw`
/// 4. Login or unlock based on login status and use that.
/// An invalid `BW_SESSION` or `--session` will prompt a login.
fn check_session_token(args: &util::Cli) -> Result<String> {
    info_println!("Getting Bitwarden session...");
    let mut session_token: String = String::new();
    // Get session flag from the user
    if args.session.trim().is_empty() {
        // No session flag, check for an environment variable
        let env_key = env::var(SESSION_ENV_KEY);
        match env_key {
            Ok(key) => {
                info_println!("{} is set. Reusing existing session.", SESSION_ENV_KEY);
                // We found it, set our session key
                session_token.push_str(&key);
            },
            Err(_) => {
                // We don't have a token to reuse, so get it from Bitwarden
                warn_println!("{} is not set. Attempting to login.", SESSION_ENV_KEY);
                let token: &String = &bwutil::get_session_token()?;

                if !token.is_empty() {
                    info_println!("Successfully unlocked. To re-use this session, run:");
                    info_println!("export {}=\"{}\"", SESSION_ENV_KEY, token);
                }
                session_token.push_str(token)
            }
        }
    }
    else {
        // Session flag is already defined, pass it on
        session_token.push_str(&args.session);
    }

    Ok(session_token)
}

/// Registers a key with `ssh-agent`.
/// First retrieves the attachment from `bitwarden-cli` and then passes it to `ssh-agent`.
fn register_key(item: &BitwardenItem, attachment_id: &str, passphrase: &str, session_token: &str) -> Result<()> {
    debug_println!("Item `{}` id({}) meets all requirements. Adding to `{}`", item.name.cyan(), item.id.to_string().cyan(), "ssh-agent".cyan());
    let bw_command = Command::new("bw")
        .arg("get")
        .arg("attachment")
        .arg(attachment_id)
        .arg("--itemid")
        .arg(&item.id)
        .arg("--raw")
        .arg("--session")
        .arg(session_token)
        .stdout(Stdio::piped())
        .spawn()
        .with_context(|| 
            format!("Error running command:\n  `bw get attachment {} --itemid {} --raw --session **REDACTED**`", attachment_id, item.id)
        )?;

    let app_path = env::current_exe()?.to_string_lossy().to_string();

    let ssh_command = Command::new("ssh-add")
        .env("SSH_ASKPASS", &app_path)
        .env(BW_FIELD_KEY_PASSPHRASE, passphrase)
        .arg("-")
        .stdin(Stdio::from(bw_command.stdout.unwrap())) // Pipe through.
        .spawn()
        .with_context(|| 
            format!("Error running command:\n  `SSH_ASKPASS=\"{}\" {}=\"{}\"ssh-add -`", app_path, BW_FIELD_KEY_PASSPHRASE, passphrase)
        )?;
    let output = ssh_command.wait_with_output().unwrap();
    let result = String::from_utf8(output.stdout).unwrap();
    if !result.is_empty() {
        println!("{:?}", result);
    }
    
    Ok(())
}