browser/apps/desktop/desktop_native/ssh_agent/examples/integration.rs

#![cfg(target_os = "linux")]

use std::{fs, process::Command, sync::Arc};

use ssh_agent::{
    agent::{
        ui_requester::{UiRequestMessage, UiRequester},
        BitwardenDesktopAgent,
    },
    memory::UnlockedSshItem,
    protocol::types::{KeyPair, PrivateKey},
    transport::unix_listener_stream::UnixListenerStream,
};
use tokio::{
    sync::{broadcast, mpsc, Mutex},
    task,
};
use tracing::info;

#[tokio::main]
async fn main() {
    let dir = homedir::my_home().unwrap().unwrap();
    let dir = dir.join(".cache");
    let dir = dir.join("ssh_agent_integration_test");
    let dir = dir.to_string_lossy().into_owned();

    // set up tracing to stdout
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .with_thread_ids(true)
        .with_thread_names(true)
        .init();

    fs::remove_dir_all(&dir).unwrap_or(());
    // Prepare test run directory
    fs::create_dir_all(&dir).unwrap();

    let config = format!(
        "Port 2222
HostKey {}/ssh_host_rsa_key
HostKey {}/ssh_host_ecdsa_key
HostKey {}/ssh_host_ed25519_key
AuthorizedKeysFile  {}/authorized_keys
",
        dir, dir, dir, dir
    );
    fs::write(format!("{}/sshd_config", dir), config).unwrap();

    let keys = make_keys(&dir);

    // Start ssh server
    let dir_clone = dir.clone();
    std::thread::spawn(move || {
        Command::new("/usr/bin/sshd")
            .args(&["-f", &format!("{}/sshd_config", &dir_clone), "-D", "-e"])
            .status()
            .expect("failed to execute process");
    });

    let ui_requester = mock_channels();
    let desktop_agent = BitwardenDesktopAgent::new(ui_requester);
    desktop_agent.set_keys(keys);

    let dir_clone = dir.clone();
    task::spawn(async move {
        println!("Starting SSH Agent V2 socket...");
        info!(target: "ssh-agent", "Listening on {}", format!("{}/ssh-agent.sock", dir_clone));
        UnixListenerStream::listen(format!("{}/ssh-agent.sock", dir_clone), desktop_agent)
            .await
            .unwrap();
    });

    // run ssh-add -L
    Command::new("ssh-add")
        .env("SSH_AUTH_SOCK", format!("{}/ssh-agent.sock", dir))
        .args(&["-L"])
        .status()
        .expect("failed to execute process");

    // run ssh
    Command::new("ssh")
        .env("SSH_AUTH_SOCK", format!("{}/ssh-agent.sock", dir))
        .args(&[
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/dev/null",
            "-p",
            "2222",
            "localhost",
            "echo",
            "Hello, world!",
        ])
        .status()
        .expect("failed to execute process");

    // Cleanup
    fs::remove_dir_all(dir).unwrap();
    std::process::exit(0);
}

fn make_keys(dir: &str) -> Vec<UnlockedSshItem> {
    Command::new("ssh-keygen")
        .args(&[
            "-f",
            &format!("{}/ssh_host_rsa_key", dir),
            "-N",
            "",
            "-t",
            "rsa",
        ])
        .status()
        .expect("failed to execute process");
    Command::new("ssh-keygen")
        .args(&[
            "-f",
            &format!("{}/ssh_host_ecdsa_key", dir),
            "-N",
            "",
            "-t",
            "ecdsa",
        ])
        .status()
        .expect("failed to execute process");
    Command::new("ssh-keygen")
        .args(&[
            "-f",
            &format!("{}/ssh_host_ed25519_key", dir),
            "-N",
            "",
            "-t",
            "ed25519",
        ])
        .status()
        .expect("failed to execute process");
    // // Make user key
    Command::new("ssh-keygen")
        .args(&[
            "-f",
            &format!("{}/id_ed25519", dir),
            "-N",
            "",
            "-t",
            "ed25519",
        ])
        .status()
        .expect("failed to execute process");
    Command::new("ssh-keygen")
        .args(&["-f", &format!("{}/ssh_rsa", dir), "-N", "", "-t", "rsa"])
        .status()
        .expect("failed to execute process");
    let pubkey1 = fs::read_to_string(format!("{}/id_ed25519.pub", dir)).unwrap();
    let pubkey2 = fs::read_to_string(format!("{}/ssh_rsa.pub", dir)).unwrap();
    fs::write(
        format!("{}/authorized_keys", dir),
        format!("{}{}", pubkey1, pubkey2),
    )
    .unwrap();
    let privkey1 = fs::read_to_string(format!("{}/id_ed25519", dir)).unwrap();
    let key1 = KeyPair::new(
        PrivateKey::try_from(privkey1).unwrap(),
        "ed25519-key".to_string(),
    );
    let privkey2 = fs::read_to_string(format!("{}/ssh_rsa", dir)).unwrap();
    let key2 = KeyPair::new(
        PrivateKey::try_from(privkey2).unwrap(),
        "rsa-key".to_string(),
    );
    let unlocked_items = vec![
        UnlockedSshItem::new(key1, "cipher1".to_string()),
        UnlockedSshItem::new(key2, "cipher2".to_string()),
    ];
    unlocked_items
}

fn mock_channels() -> UiRequester {
    let (show_ui_request_tx, mut show_ui_request_rx) = mpsc::channel::<UiRequestMessage>(10);

    // Create mock broadcast channel for responses
    let (response_tx, response_rx) = broadcast::channel::<(u32, bool)>(10);
    let get_ui_response_rx = Arc::new(Mutex::new(response_rx));

    // Spawn a task to automatically send back "true" responses
    let response_tx_clone = response_tx.clone();
    let _ = task::spawn(async move {
        while let Some(req) = show_ui_request_rx.recv().await {
            info!("Mock UI requester received request: {:?}", req);
            let _ = response_tx_clone.send((req.id(), true));
        }
        info!("Mock UI requester task ending");
    });

    UiRequester::new(show_ui_request_tx, get_ui_response_rx)
}