sal-modular/wasm_app/src/vault_bindings.rs

//! WebAssembly bindings for accessing vault operations (session, keypairs, signing, scripting, etc)
#![cfg(target_arch = "wasm32")]

use kvstore::wasm::WasmStore;
use once_cell::unsync::Lazy;
use rhai::Engine;
use std::cell::RefCell;
use vault::rhai_bindings as vault_rhai_bindings;
use vault::session::SessionManager;
use wasm_bindgen::prelude::*;
use wasm_bindgen::JsValue;
use js_sys::Uint8Array;

thread_local! {
    static ENGINE: Lazy<RefCell<Engine>> = Lazy::new(|| RefCell::new(Engine::new()));
    static SESSION_PASSWORD: RefCell<Option<Vec<u8>>> = RefCell::new(None);
}

pub use vault::session_singleton::SESSION_MANAGER;

// =====================
// Session Lifecycle
// =====================

/// Create and unlock a new keyspace with the given name and password
#[wasm_bindgen]
pub async fn create_keyspace(keyspace: &str, password: &str) -> Result<(), JsValue> {
    let keyspace = keyspace.to_string();
    let password_vec = password.as_bytes().to_vec();
    match WasmStore::open("vault").await {
        Ok(store) => {
            let vault = vault::Vault::new(store);
            let mut manager = SessionManager::new(vault);
            match manager.create_keyspace(&keyspace, &password_vec, None).await {
                Ok(_) => {
                    SESSION_MANAGER.with(|cell| cell.replace(Some(manager)));
                }
                Err(e) => {
                    web_sys::console::error_1(&format!("Failed to create keyspace: {e}").into());
                    return Err(JsValue::from_str(&format!("Failed to create keyspace: {e}")));
                }
            }
        }
        Err(e) => {
            web_sys::console::error_1(&format!("Failed to open WasmStore: {e}").into());
            return Err(JsValue::from_str(&format!("Failed to open WasmStore: {e}")));
        }
    }
    SESSION_PASSWORD.with(|cell| cell.replace(Some(password.as_bytes().to_vec())));
    Ok(())
}

/// Initialize session with keyspace and password
#[wasm_bindgen]
pub async fn init_session(keyspace: &str, password: &str) -> Result<(), JsValue> {
    let keyspace = keyspace.to_string();
    let password_vec = password.as_bytes().to_vec();
    match WasmStore::open("vault").await {
        Ok(store) => {
            let vault = vault::Vault::new(store);
            let mut manager = SessionManager::new(vault);
            match manager.unlock_keyspace(&keyspace, &password_vec).await {
                Ok(_) => {
                    SESSION_MANAGER.with(|cell| cell.replace(Some(manager)));
                }
                Err(e) => {
                    web_sys::console::error_1(&format!("Failed to unlock keyspace: {e}").into());
                    return Err(JsValue::from_str(&format!("Failed to unlock keyspace: {e}")));
                }
            }
        }
        Err(e) => {
            web_sys::console::error_1(&format!("Failed to open WasmStore: {e}").into());
            return Err(JsValue::from_str(&format!("Failed to open WasmStore: {e}")));
        }
    }
    SESSION_PASSWORD.with(|cell| cell.replace(Some(password.as_bytes().to_vec())));
    Ok(())
}


/// Lock the session (zeroize password and session)
#[wasm_bindgen]
pub fn lock_session() {
    SESSION_MANAGER.with(|cell| *cell.borrow_mut() = None);
    SESSION_PASSWORD.with(|cell| *cell.borrow_mut() = None);
}

// =====================
// Keypair Management
// =====================

/// Get metadata of the currently selected keypair
#[wasm_bindgen]
pub fn current_keypair_metadata() -> Result<JsValue, JsValue> {
    SESSION_MANAGER.with(|cell| {
        cell.borrow().as_ref()
            .and_then(|session| session.current_keypair_metadata())
            .map(|meta| wasm_bindgen::JsValue::from_serde(&meta).unwrap())
            .ok_or_else(|| JsValue::from_str("No keypair selected or no keyspace unlocked"))
    })
}

/// Get public key of the currently selected keypair as Uint8Array
#[wasm_bindgen]
pub fn current_keypair_public_key() -> Result<JsValue, JsValue> {
    SESSION_MANAGER.with(|cell| {
        cell.borrow().as_ref()
            .and_then(|session| session.current_keypair_public_key())
            .map(|pk| js_sys::Uint8Array::from(pk.as_slice()).into())
            .ok_or_else(|| JsValue::from_str("No keypair selected or no keyspace unlocked"))
    })
}

/// Returns true if a keyspace is currently unlocked
#[wasm_bindgen]
pub fn is_unlocked() -> bool {
    SESSION_MANAGER.with(|cell| {
        cell.borrow().as_ref().map(|session| session.is_unlocked()).unwrap_or(false)
    })
}

/// Get the default public key for a workspace (keyspace)
/// This returns the public key of the first keypair in the keyspace
#[wasm_bindgen]
pub async fn get_workspace_default_public_key(workspace_id: &str) -> Result<JsValue, JsValue> {
    // For now, workspace_id is the same as keyspace name
    // In a full implementation, you might have a mapping from workspace to keyspace

    SESSION_MANAGER.with(|cell| {
        if let Some(session) = cell.borrow().as_ref() {
            if let Some(keyspace_name) = session.current_keyspace_name() {
                if keyspace_name == workspace_id {
                    // Use the default_keypair method to get the first keypair
                    if let Some(default_keypair) = session.default_keypair() {
                        // Return the actual public key as hex
                        let public_key_hex = hex::encode(&default_keypair.public_key);
                        return Ok(JsValue::from_str(&public_key_hex));
                    }
                }
            }
        }
        Err(JsValue::from_str("Workspace not found or no keypairs available"))
    })
}

/// Get the current unlocked public key as hex string
#[wasm_bindgen]
pub fn get_current_unlocked_public_key() -> Result<String, JsValue> {
    SESSION_MANAGER.with(|cell| {
        cell.borrow().as_ref()
            .and_then(|session| session.current_keypair_public_key())
            .map(|pk| hex::encode(pk.as_slice()))
            .ok_or_else(|| JsValue::from_str("No keypair selected or no keyspace unlocked"))
    })
}

/// Get all keypairs from the current session
/// Returns an array of keypair objects with id, type, and metadata
// #[wasm_bindgen]
// pub async fn list_keypairs() -> Result<JsValue, JsValue> {
//     // [Function body commented out to resolve duplicate symbol error]
//     // (Original implementation moved to keypair_bindings.rs)
//     unreachable!("This function is disabled. Use the export from keypair_bindings.rs.");
// }

// [Function body commented out to resolve duplicate symbol error]
// }


/// Select keypair for the session
#[wasm_bindgen]
pub fn select_keypair(key_id: &str) -> Result<(), JsValue> {
    let mut result = Err(JsValue::from_str("Session not initialized"));
    SESSION_MANAGER.with(|cell| {
        if let Some(session) = cell.borrow_mut().as_mut() {
            result = session
                .select_keypair(key_id)
                .map_err(|e| JsValue::from_str(&format!("select_keypair error: {e}")));
        }
    });
    result
}


/// List keypairs in the current session's keyspace
#[wasm_bindgen]
pub async fn list_keypairs() -> Result<JsValue, JsValue> {
    SESSION_MANAGER.with(|cell| {
        if let Some(session) = cell.borrow().as_ref() {
            if let Some(keyspace) = session.current_keyspace() {
                let keypairs = &keyspace.keypairs;
                serde_json::to_string(keypairs)
                    .map(|s| JsValue::from_str(&s))
                    .map_err(|e| JsValue::from_str(&format!("Serialization error: {e}")))
            } else {
                Err(JsValue::from_str("No keyspace unlocked"))
            }
        } else {
            Err(JsValue::from_str("Session not initialized"))
        }
    })
}

/// Add a keypair to the current keyspace
#[wasm_bindgen]
pub async fn add_keypair(
    key_type: Option<String>,
    metadata: Option<String>,
) -> Result<JsValue, JsValue> {
    use vault::{KeyMetadata, KeyType};
    let password = SESSION_PASSWORD
        .with(|pw| pw.borrow().clone())
        .ok_or_else(|| JsValue::from_str("Session password not set"))?;
    let keyspace_name = SESSION_MANAGER.with(|cell| {
        cell.borrow().as_ref().and_then(|session| {
            session.current_keyspace_name().map(|name| name.to_string())
        })
    });
    let keyspace_name = keyspace_name.ok_or_else(|| JsValue::from_str("No keyspace selected"))?;
    let key_type = key_type
        .as_deref()
        .map(|s| match s {
            "Ed25519" => KeyType::Ed25519,
            "Secp256k1" => KeyType::Secp256k1,
            _ => KeyType::Secp256k1,
        })
        .unwrap_or(KeyType::Secp256k1);
    let metadata = match metadata {
        Some(ref meta_str) => Some(
            serde_json::from_str::<KeyMetadata>(meta_str)
                .map_err(|e| JsValue::from_str(&format!("Invalid metadata: {e}")))?,
        ),
        None => None,
    };
    // Take session out, do async work, then put it back
    let mut session_opt = SESSION_MANAGER.with(|cell| cell.borrow_mut().take());
    let session = session_opt.as_mut().ok_or_else(|| JsValue::from_str("Session not initialized"))?;
    let key_id = session
        .get_vault_mut()
        .add_keypair(&keyspace_name, &password, Some(key_type), metadata)
        .await
        .map_err(|e| JsValue::from_str(&format!("add_keypair error: {e}")))?;
    // Refresh in-memory keyspace data so list_keypairs reflects the new keypair immediately
    session.unlock_keyspace(&keyspace_name, &password).await
        .map_err(|e| JsValue::from_str(&format!("refresh keyspace after add_keypair error: {e}")))?;
    // Put session back
    SESSION_MANAGER.with(|cell| *cell.borrow_mut() = Some(session_opt.take().unwrap()));
    Ok(JsValue::from_str(&key_id))
}

/// Sign message with current session (requires selected keypair)
#[wasm_bindgen]
pub async fn sign(message: &[u8]) -> Result<JsValue, JsValue> {
    {
        // SAFETY: We only use this pointer synchronously within this function, and SESSION_MANAGER outlives this scope.
        let session_ptr =
            SESSION_MANAGER.with(|cell| cell.borrow().as_ref().map(|s| s as *const _));
        let session: &vault::session::SessionManager<kvstore::wasm::WasmStore> = match session_ptr {
            Some(ptr) => unsafe { &*ptr },
            None => return Err(JsValue::from_str("Session not initialized")),
        };
        match session.sign(message).await {
            Ok(sig_bytes) => {
                let hex_sig = hex::encode(&sig_bytes);
                Ok(JsValue::from_str(&hex_sig))
            }
            Err(e) => Err(JsValue::from_str(&format!("Sign error: {e}"))),
        }
    }
}

/// Sign message with default keypair (first keypair in keyspace) without changing session state
#[wasm_bindgen]
pub async fn sign_with_default_keypair(message: &[u8]) -> Result<JsValue, JsValue> {
    // Temporarily select the default keypair, sign, then restore the original selection
    let original_keypair = SESSION_MANAGER.with(|cell| {
        cell.borrow().as_ref()
            .and_then(|session| session.current_keypair())
            .map(|kp| kp.id.clone())
    });

    // Select default keypair
    let select_result = SESSION_MANAGER.with(|cell| {
        let mut session_opt = cell.borrow_mut().take();
        if let Some(ref mut session) = session_opt {
            let result = session.select_default_keypair();
            *cell.borrow_mut() = Some(session_opt.take().unwrap());
            result.map_err(|e| e.to_string())
        } else {
            Err("Session not initialized".to_string())
        }
    });

    if let Err(e) = select_result {
        return Err(JsValue::from_str(&format!("Failed to select default keypair: {e}")));
    }

    // Sign with the default keypair
    let sign_result = {
        let session_ptr = SESSION_MANAGER.with(|cell| cell.borrow().as_ref().map(|s| s as *const _));
        let session: &vault::session::SessionManager<kvstore::wasm::WasmStore> = match session_ptr {
            Some(ptr) => unsafe { &*ptr },
            None => return Err(JsValue::from_str("Session not initialized")),
        };
        session.sign(message).await
    };

    // Restore original keypair selection if there was one
    if let Some(original_id) = original_keypair {
        SESSION_MANAGER.with(|cell| {
            let mut session_opt = cell.borrow_mut().take();
            if let Some(ref mut session) = session_opt {
                let _ = session.select_keypair(&original_id); // Ignore errors here
                *cell.borrow_mut() = Some(session_opt.take().unwrap());
            }
        });
    }

    // Return the signature result
    match sign_result {
        Ok(sig_bytes) => {
            let hex_sig = hex::encode(&sig_bytes);
            Ok(JsValue::from_str(&hex_sig))
        }
        Err(e) => Err(JsValue::from_str(&format!("Sign error: {e}"))),
    }
}

/// Verify a signature with the current session's selected keypair
#[wasm_bindgen]
pub async fn verify(message: &[u8], signature: &str) -> Result<JsValue, JsValue> {
    {
        // SAFETY: We only use this pointer synchronously within this function, and SESSION_MANAGER outlives this scope.
        let session_ptr =
            SESSION_MANAGER.with(|cell| cell.borrow().as_ref().map(|s| s as *const _));
        let session: &vault::session::SessionManager<kvstore::wasm::WasmStore> = match session_ptr {
            Some(ptr) => unsafe { &*ptr },
            None => return Err(JsValue::from_str("Session not initialized")),
        };
        
        // Convert hex signature to bytes
        let sig_bytes = match hex::decode(signature) {
            Ok(bytes) => bytes,
            Err(e) => return Err(JsValue::from_str(&format!("Invalid signature format: {e}"))),
        };
        
        match session.verify(message, &sig_bytes).await {
            Ok(is_valid) => Ok(JsValue::from_bool(is_valid)),
            Err(e) => Err(JsValue::from_str(&format!("Verify error: {e}"))),
        }
    }
}

/// Encrypt data using the current session's keyspace symmetric cipher
#[wasm_bindgen]
pub async fn encrypt_data(data: &[u8]) -> Result<JsValue, JsValue> {
    {
        // SAFETY: We only use this pointer synchronously within this function, and SESSION_MANAGER outlives this scope.
        let session_ptr =
            SESSION_MANAGER.with(|cell| cell.borrow().as_ref().map(|s| s as *const _));
        let session: &vault::session::SessionManager<kvstore::wasm::WasmStore> = match session_ptr {
            Some(ptr) => unsafe { &*ptr },
            None => return Err(JsValue::from_str("Session not initialized")),
        };
        
        match session.encrypt(data).await {
            Ok(encrypted) => {
                // Return as Uint8Array for JavaScript
                Ok(Uint8Array::from(&encrypted[..]).into())
            }
            Err(e) => Err(JsValue::from_str(&format!("Encryption error: {e}"))),
        }
    }
}

/// Decrypt data using the current session's keyspace symmetric cipher
#[wasm_bindgen]
pub async fn decrypt_data(encrypted: &[u8]) -> Result<JsValue, JsValue> {
    {
        // SAFETY: We only use this pointer synchronously within this function, and SESSION_MANAGER outlives this scope.
        let session_ptr =
            SESSION_MANAGER.with(|cell| cell.borrow().as_ref().map(|s| s as *const _));
        let session: &vault::session::SessionManager<kvstore::wasm::WasmStore> = match session_ptr {
            Some(ptr) => unsafe { &*ptr },
            None => return Err(JsValue::from_str("Session not initialized")),
        };
        
        match session.decrypt(encrypted).await {
            Ok(decrypted) => {
                // Return as Uint8Array for JavaScript
                Ok(Uint8Array::from(&decrypted[..]).into())
            }
            Err(e) => Err(JsValue::from_str(&format!("Decryption error: {e}"))),
        }
    }
}