index : bitwarden-ssh-agent.git

use std::{env, process::{Command, Stdio}};
use gumdrop::Options;
use anyhow::{Result, Context};

mod util;
mod bwutil;

/// Environment variable housing an existing Bitwarden session.
const SESSION_ENV_KEY: &str = "BW_SESSION";
/// 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 {
    ($debug:expr, $($arg:tt)*) => {
		if $debug {
			println!($($arg)*);
		}
    }
}

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(())
	}

    let args: util::Cli = util::Cli::parse_args_default_or_exit();
	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!(args.debug, "Found {} folder(s) named `{}`", folders.len(), args.folder);
		
		// Retrieve items from each folder since there may be multiple folders with the same name.
		for folder in folders {
			let folder_items = bwutil::exec_list_folder_items(&session_token, &folder.id)?;
			debug_println!(args.debug, "Found {} item(s) in folder `{}` id({})", folder_items.len(), folder.name, folder.id);
			
			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 mut key_filename = String::new();
					let mut key_passphrase = String::new();
					for field in fields {
						if field.name.to_lowercase().eq(&args.key.to_lowercase()) {
							key_filename.push_str(&field.value);
						}
						if field.name.to_lowercase().eq(&args.passphrase.to_lowercase()) {
							key_passphrase.push_str(&field.value);
						}
					}

					if let Some(attachments) = &item.attachments {
						for attachment in attachments {
							if attachment.file_name.eq(&key_filename) {
								debug_println!(args.debug, "Item `{}` id({}) meets all requirements. Adding to `ssh-agent`", item.name, item.id);
								let _key = register_key(&item.id, &attachment.id, &key_passphrase, &session_token)?;
							}
						}
					}
				}
			}
		}
	}

	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> {
	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) => {
				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
				println!("{} is not set. Attempting to login.", SESSION_ENV_KEY);
				let token: &String = &bwutil::get_session_token()?;

				if !token.is_empty() {
					println!("Successfully unlocked. To re-use this session, run:");
					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)
}

fn register_key(item_id: &str, attachment_id: &str, passphrase: &str, session_token: &str) -> Result<()> {
    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()?;
	let app_path = app_path.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(())
}