From 3a0900fc15c4a6a5593260bbf3b3ef0d9636d555 Mon Sep 17 00:00:00 2001 From: Mahmoud Emad Date: Mon, 12 May 2025 12:47:37 +0300 Subject: [PATCH] refactor: Improve Rhai test runner and vault module code - Updated the Rhai test runner script to correctly find test files. - Improved the structure and formatting of the `vault.rs` module. - Minor code style improvements in multiple files. --- run_rhai_tests.sh | 2 +- src/rhai/vault.rs | 179 +++++++++++++-------------- src/vault/keypair/session_manager.rs | 49 ++++---- src/vault/kvs/error.rs | 19 +-- src/vault/kvs/store.rs | 46 ++++--- 5 files changed, 154 insertions(+), 141 deletions(-) diff --git a/run_rhai_tests.sh b/run_rhai_tests.sh index 2182cb5..4b7fb08 100755 --- a/run_rhai_tests.sh +++ b/run_rhai_tests.sh @@ -24,7 +24,7 @@ log "${BLUE} Running All Rhai Tests ${NC}" log "${BLUE}=======================================${NC}" # Find all test runner scripts -RUNNERS=$(find src/rhai_tests -name "run_all_tests.rhai") +RUNNERS=$(find rhai_tests -name "run_all_tests.rhai") # Initialize counters TOTAL_MODULES=0 diff --git a/src/rhai/vault.rs b/src/rhai/vault.rs index 3f703ee..42f68d4 100644 --- a/src/rhai/vault.rs +++ b/src/rhai/vault.rs @@ -1,29 +1,28 @@ //! Rhai bindings for SAL crypto functionality -use rhai::{Engine, Dynamic, EvalAltResult}; -use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; +use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _}; +use ethers::types::{Address, U256}; +use once_cell::sync::Lazy; +use rhai::{Dynamic, Engine, EvalAltResult}; +use std::collections::HashMap; use std::fs; use std::path::PathBuf; -use std::collections::HashMap; -use std::sync::Mutex; -use once_cell::sync::Lazy; -use tokio::runtime::Runtime; -use ethers::types::{Address, U256}; use std::str::FromStr; +use std::sync::Mutex; +use tokio::runtime::Runtime; -use crate::vault::{keypair, ethereum}; -use crate::vault::ethereum::contract_utils::{prepare_function_arguments, convert_token_to_rhai}; +use crate::vault::ethereum::contract_utils::{convert_token_to_rhai, prepare_function_arguments}; +use crate::vault::{ethereum, keypair}; -use symmetric_impl::implementation as symmetric_impl; +use crate::vault::symmetric::implementation as symmetric_impl; // Global Tokio runtime for blocking async operations -static RUNTIME: Lazy> = Lazy::new(|| { - Mutex::new(Runtime::new().expect("Failed to create Tokio runtime")) -}); +static RUNTIME: Lazy> = + Lazy::new(|| Mutex::new(Runtime::new().expect("Failed to create Tokio runtime"))); // Global provider registry -static PROVIDERS: Lazy>>> = Lazy::new(|| { - Mutex::new(HashMap::new()) -}); +static PROVIDERS: Lazy< + Mutex>>, +> = Lazy::new(|| Mutex::new(HashMap::new())); // Key space management functions fn load_key_space(name: &str, password: &str) -> bool { @@ -90,7 +89,8 @@ fn create_key_space(name: &str, password: &str) -> bool { match keypair::get_current_space() { Ok(space) => { // Encrypt the key space - let encrypted_space = match symmetric_impl::encrypt_key_space(&space, password) { + let encrypted_space = match symmetric_impl::encrypt_key_space(&space, password) + { Ok(encrypted) => encrypted, Err(e) => { log::error!("Error encrypting key space: {}", e); @@ -99,13 +99,14 @@ fn create_key_space(name: &str, password: &str) -> bool { }; // Serialize the encrypted space - let serialized = match symmetric_impl::serialize_encrypted_space(&encrypted_space) { - Ok(json) => json, - Err(e) => { - log::error!("Error serializing encrypted space: {}", e); - return false; - } - }; + let serialized = + match symmetric_impl::serialize_encrypted_space(&encrypted_space) { + Ok(json) => json, + Err(e) => { + log::error!("Error serializing encrypted space: {}", e); + return false; + } + }; // Get the key spaces directory let home_dir = dirs::home_dir().unwrap_or_else(|| PathBuf::from(".")); @@ -114,7 +115,7 @@ fn create_key_space(name: &str, password: &str) -> bool { // Create directory if it doesn't exist if !key_spaces_dir.exists() { match fs::create_dir_all(&key_spaces_dir) { - Ok(_) => {}, + Ok(_) => {} Err(e) => { log::error!("Error creating key spaces directory: {}", e); return false; @@ -128,19 +129,19 @@ fn create_key_space(name: &str, password: &str) -> bool { Ok(_) => { log::info!("Key space created and saved to {}", space_path.display()); true - }, + } Err(e) => { log::error!("Error writing key space file: {}", e); false } } - }, + } Err(e) => { log::error!("Error getting current space: {}", e); false } } - }, + } Err(e) => { log::error!("Error creating key space: {}", e); false @@ -177,7 +178,7 @@ fn auto_save_key_space(password: &str) -> bool { // Create directory if it doesn't exist if !key_spaces_dir.exists() { match fs::create_dir_all(&key_spaces_dir) { - Ok(_) => {}, + Ok(_) => {} Err(e) => { log::error!("Error creating key spaces directory: {}", e); return false; @@ -191,13 +192,13 @@ fn auto_save_key_space(password: &str) -> bool { Ok(_) => { log::info!("Key space saved to {}", space_path.display()); true - }, + } Err(e) => { log::error!("Error writing key space file: {}", e); false } } - }, + } Err(e) => { log::error!("Error getting current space: {}", e); false @@ -207,21 +208,17 @@ fn auto_save_key_space(password: &str) -> bool { fn encrypt_key_space(password: &str) -> String { match keypair::get_current_space() { - Ok(space) => { - match symmetric_impl::encrypt_key_space(&space, password) { - Ok(encrypted_space) => { - match serde_json::to_string(&encrypted_space) { - Ok(json) => json, - Err(e) => { - log::error!("Error serializing encrypted space: {}", e); - String::new() - } - } - }, + Ok(space) => match symmetric_impl::encrypt_key_space(&space, password) { + Ok(encrypted_space) => match serde_json::to_string(&encrypted_space) { + Ok(json) => json, Err(e) => { - log::error!("Error encrypting key space: {}", e); + log::error!("Error serializing encrypted space: {}", e); String::new() } + }, + Err(e) => { + log::error!("Error encrypting key space: {}", e); + String::new() } }, Err(e) => { @@ -235,13 +232,11 @@ fn decrypt_key_space(encrypted: &str, password: &str) -> bool { match serde_json::from_str(encrypted) { Ok(encrypted_space) => { match symmetric_impl::decrypt_key_space(&encrypted_space, password) { - Ok(space) => { - match keypair::set_current_space(space) { - Ok(_) => true, - Err(e) => { - log::error!("Error setting current space: {}", e); - false - } + Ok(space) => match keypair::set_current_space(space) { + Ok(_) => true, + Err(e) => { + log::error!("Error setting current space: {}", e); + false } }, Err(e) => { @@ -249,7 +244,7 @@ fn decrypt_key_space(encrypted: &str, password: &str) -> bool { false } } - }, + } Err(e) => { log::error!("Error parsing encrypted space: {}", e); false @@ -263,7 +258,7 @@ fn create_keypair(name: &str, password: &str) -> bool { Ok(_) => { // Auto-save the key space after creating a keypair auto_save_key_space(password) - }, + } Err(e) => { log::error!("Error creating keypair: {}", e); false @@ -306,13 +301,11 @@ fn sign(message: &str) -> String { fn verify(message: &str, signature: &str) -> bool { let message_bytes = message.as_bytes(); match BASE64.decode(signature) { - Ok(signature_bytes) => { - match keypair::keypair_verify(message_bytes, &signature_bytes) { - Ok(is_valid) => is_valid, - Err(e) => { - log::error!("Error verifying signature: {}", e); - false - } + Ok(signature_bytes) => match keypair::keypair_verify(message_bytes, &signature_bytes) { + Ok(is_valid) => is_valid, + Err(e) => { + log::error!("Error verifying signature: {}", e); + false } }, Err(e) => { @@ -339,7 +332,7 @@ fn encrypt(key: &str, message: &str) -> String { String::new() } } - }, + } Err(e) => { log::error!("Error decoding key: {}", e); String::new() @@ -349,30 +342,26 @@ fn encrypt(key: &str, message: &str) -> String { fn decrypt(key: &str, ciphertext: &str) -> String { match BASE64.decode(key) { - Ok(key_bytes) => { - match BASE64.decode(ciphertext) { - Ok(ciphertext_bytes) => { - match symmetric_impl::decrypt_symmetric(&key_bytes, &ciphertext_bytes) { - Ok(plaintext) => { - match String::from_utf8(plaintext) { - Ok(text) => text, - Err(e) => { - log::error!("Error converting plaintext to string: {}", e); - String::new() - } - } - }, + Ok(key_bytes) => match BASE64.decode(ciphertext) { + Ok(ciphertext_bytes) => { + match symmetric_impl::decrypt_symmetric(&key_bytes, &ciphertext_bytes) { + Ok(plaintext) => match String::from_utf8(plaintext) { + Ok(text) => text, Err(e) => { - log::error!("Error decrypting ciphertext: {}", e); + log::error!("Error converting plaintext to string: {}", e); String::new() } + }, + Err(e) => { + log::error!("Error decrypting ciphertext: {}", e); + String::new() } - }, - Err(e) => { - log::error!("Error decoding ciphertext: {}", e); - String::new() } } + Err(e) => { + log::error!("Error decoding ciphertext: {}", e); + String::new() + } }, Err(e) => { log::error!("Error decoding key: {}", e); @@ -478,7 +467,11 @@ fn get_wallet_address_for_network(network_name: &str) -> String { match ethereum::get_current_ethereum_wallet_for_network(network_name_proper) { Ok(wallet) => wallet.address_string(), Err(e) => { - log::error!("Error getting wallet address for network {}: {}", network_name, e); + log::error!( + "Error getting wallet address for network {}: {}", + network_name, + e + ); String::new() } } @@ -542,7 +535,11 @@ fn create_wallet_from_private_key_for_network(private_key: &str, network_name: & match ethereum::create_ethereum_wallet_from_private_key_for_network(private_key, network) { Ok(_) => true, Err(e) => { - log::error!("Error creating wallet from private key for network {}: {}", network_name, e); + log::error!( + "Error creating wallet from private key for network {}: {}", + network_name, + e + ); false } } @@ -563,7 +560,7 @@ fn create_agung_provider() -> String { log::error!("Failed to acquire provider registry lock"); String::new() - }, + } Err(e) => { log::error!("Error creating Agung provider: {}", e); String::new() @@ -619,9 +616,7 @@ fn get_balance(network_name: &str, address: &str) -> String { }; // Execute the balance query in a blocking manner - match rt.block_on(async { - ethereum::get_balance(&provider, addr).await - }) { + match rt.block_on(async { ethereum::get_balance(&provider, addr).await }) { Ok(balance) => balance.to_string(), Err(e) => { log::error!("Failed to get balance: {}", e); @@ -687,9 +682,7 @@ fn send_eth(wallet_network: &str, to_address: &str, amount_str: &str) -> String }; // Execute the transaction in a blocking manner - match rt.block_on(async { - ethereum::send_eth(&wallet, &provider, to_addr, amount).await - }) { + match rt.block_on(async { ethereum::send_eth(&wallet, &provider, to_addr, amount).await }) { Ok(tx_hash) => format!("{:?}", tx_hash), Err(e) => { log::error!("Transaction failed: {}", e); @@ -731,7 +724,7 @@ fn load_contract_abi(network_name: &str, address: &str, abi_json: &str) -> Strin String::new() } } - }, + } Err(e) => { log::error!("Error creating contract: {}", e); String::new() @@ -916,14 +909,20 @@ pub fn register_crypto_module(engine: &mut Engine) -> Result<(), Box> = Lazy::new(|| { - Mutex::new(Session::default()) -}); +pub static SESSION: Lazy> = Lazy::new(|| Mutex::new(Session::default())); // Session management and selected keypair operation functions will be added here /// Creates a new key space with the given name. pub fn create_space(name: &str) -> Result<(), CryptoError> { let mut session = SESSION.lock().unwrap(); - + // Create a new space let space = KeySpace::new(name); - + // Set as current space session.current_space = Some(space); session.selected_keypair = None; - + Ok(()) } @@ -53,7 +48,10 @@ pub fn set_current_space(space: KeySpace) -> Result<(), CryptoError> { /// Gets the current key space. pub fn get_current_space() -> Result { let session = SESSION.lock().unwrap(); - session.current_space.clone().ok_or(CryptoError::NoActiveSpace) + session + .current_space + .clone() + .ok_or(CryptoError::NoActiveSpace) } /// Clears the current session (logout). @@ -66,18 +64,18 @@ pub fn clear_session() { /// Creates a new keypair in the current space. pub fn create_keypair(name: &str) -> Result<(), CryptoError> { let mut session = SESSION.lock().unwrap(); - + if let Some(ref mut space) = session.current_space { if space.keypairs.contains_key(name) { return Err(CryptoError::KeypairAlreadyExists(name.to_string())); } - + let keypair = KeyPair::new(name); space.keypairs.insert(name.to_string(), keypair); - + // Automatically select the new keypair session.selected_keypair = Some(name.to_string()); - + Ok(()) } else { Err(CryptoError::NoActiveSpace) @@ -87,12 +85,12 @@ pub fn create_keypair(name: &str) -> Result<(), CryptoError> { /// Selects a keypair for use. pub fn select_keypair(name: &str) -> Result<(), CryptoError> { let mut session = SESSION.lock().unwrap(); - + if let Some(ref space) = session.current_space { if !space.keypairs.contains_key(name) { return Err(CryptoError::KeypairNotFound(name.to_string())); } - + session.selected_keypair = Some(name.to_string()); Ok(()) } else { @@ -103,7 +101,7 @@ pub fn select_keypair(name: &str) -> Result<(), CryptoError> { /// Gets the currently selected keypair. pub fn get_selected_keypair() -> Result { let session = SESSION.lock().unwrap(); - + if let Some(ref space) = session.current_space { if let Some(ref keypair_name) = session.selected_keypair { if let Some(keypair) = space.keypairs.get(keypair_name) { @@ -113,14 +111,14 @@ pub fn get_selected_keypair() -> Result { } return Err(CryptoError::NoKeypairSelected); } - + Err(CryptoError::NoActiveSpace) } /// Lists all keypair names in the current space. pub fn list_keypairs() -> Result, CryptoError> { let session = SESSION.lock().unwrap(); - + if let Some(ref space) = session.current_space { Ok(space.keypairs.keys().cloned().collect()) } else { @@ -152,12 +150,19 @@ pub fn keypair_verify(message: &[u8], signature_bytes: &[u8]) -> Result Result { +pub fn verify_with_public_key( + public_key: &[u8], + message: &[u8], + signature_bytes: &[u8], +) -> Result { KeyPair::verify_with_public_key(public_key, message, signature_bytes) } /// Encrypts a message for a recipient using their public key. -pub fn encrypt_asymmetric(recipient_public_key: &[u8], message: &[u8]) -> Result, CryptoError> { +pub fn encrypt_asymmetric( + recipient_public_key: &[u8], + message: &[u8], +) -> Result, CryptoError> { let keypair = get_selected_keypair()?; keypair.encrypt_asymmetric(recipient_public_key, message) } @@ -166,4 +171,4 @@ pub fn encrypt_asymmetric(recipient_public_key: &[u8], message: &[u8]) -> Result pub fn decrypt_asymmetric(ciphertext: &[u8]) -> Result, CryptoError> { let keypair = get_selected_keypair()?; keypair.decrypt_asymmetric(ciphertext) -} \ No newline at end of file +} diff --git a/src/vault/kvs/error.rs b/src/vault/kvs/error.rs index 5fcd15c..bbd6eaf 100644 --- a/src/vault/kvs/error.rs +++ b/src/vault/kvs/error.rs @@ -1,6 +1,5 @@ //! Error types for the key-value store. -use std::fmt; use thiserror::Error; /// Errors that can occur when using the key-value store. @@ -9,31 +8,31 @@ pub enum KvsError { /// I/O error #[error("I/O error: {0}")] Io(#[from] std::io::Error), - + /// Key not found #[error("Key not found: {0}")] KeyNotFound(String), - + /// Store not found #[error("Store not found: {0}")] StoreNotFound(String), - + /// Serialization error #[error("Serialization error: {0}")] Serialization(String), - + /// Deserialization error #[error("Deserialization error: {0}")] Deserialization(String), - + /// Encryption error #[error("Encryption error: {0}")] Encryption(String), - + /// Decryption error #[error("Decryption error: {0}")] Decryption(String), - + /// Other error #[error("Error: {0}")] Other(String), @@ -56,7 +55,9 @@ impl From for KvsError { match err { crate::vault::error::CryptoError::EncryptionFailed(msg) => KvsError::Encryption(msg), crate::vault::error::CryptoError::DecryptionFailed(msg) => KvsError::Decryption(msg), - crate::vault::error::CryptoError::SerializationError(msg) => KvsError::Serialization(msg), + crate::vault::error::CryptoError::SerializationError(msg) => { + KvsError::Serialization(msg) + } _ => KvsError::Other(err.to_string()), } } diff --git a/src/vault/kvs/store.rs b/src/vault/kvs/store.rs index 775fcdc..74c9c6f 100644 --- a/src/vault/kvs/store.rs +++ b/src/vault/kvs/store.rs @@ -1,7 +1,9 @@ //! Implementation of a simple key-value store using the filesystem. use crate::vault::kvs::error::{KvsError, Result}; -use crate::vault::symmetric::implementation::{derive_key_from_password, encrypt_symmetric, decrypt_symmetric}; +use crate::vault::symmetric::implementation::{ + decrypt_symmetric, derive_key_from_password, encrypt_symmetric, +}; use serde::{de::DeserializeOwned, Deserialize, Serialize}; use std::collections::HashMap; use std::fs; @@ -52,7 +54,9 @@ pub fn get_store_path() -> PathBuf { pub fn create_store(name: &str, encrypted: bool, password: Option<&str>) -> Result { // Check if password is provided when encryption is enabled if encrypted && password.is_none() { - return Err(KvsError::Other("Password required for encrypted store".to_string())); + return Err(KvsError::Other( + "Password required for encrypted store".to_string(), + )); } // Create the store directory if it doesn't exist @@ -107,7 +111,9 @@ pub fn open_store(name: &str, password: Option<&str>) -> Result { // If encrypted, we need a password if is_encrypted && password.is_none() { - return Err(KvsError::Other("Password required for encrypted store".to_string())); + return Err(KvsError::Other( + "Password required for encrypted store".to_string(), + )); } // Parse the store data @@ -115,8 +121,8 @@ pub fn open_store(name: &str, password: Option<&str>) -> Result { // Decrypt the file content let password = password.unwrap(); let encrypted_data: Vec = serde_json::from_str(&file_content)?; - let key = implementation::derive_key_from_password(password); - let decrypted_data = implementation::decrypt_symmetric(&key, &encrypted_data)?; + let key = derive_key_from_password(password); + let decrypted_data = decrypt_symmetric(&key, &encrypted_data)?; let decrypted_str = String::from_utf8(decrypted_data) .map_err(|e| KvsError::Deserialization(e.to_string()))?; serde_json::from_str(&decrypted_str)? @@ -203,12 +209,14 @@ impl KvStore { if self.encrypted { if let Some(password) = &self.password { // Encrypt the data - let key = implementation::derive_key_from_password(password); - let encrypted_data = implementation::encrypt_symmetric(&key, serialized.as_bytes())?; + let key = derive_key_from_password(password); + let encrypted_data = encrypt_symmetric(&key, serialized.as_bytes())?; let encrypted_json = serde_json::to_string(&encrypted_data)?; fs::write(&self.path, encrypted_json)?; } else { - return Err(KvsError::Other("Password required for encrypted store".to_string())); + return Err(KvsError::Other( + "Password required for encrypted store".to_string(), + )); } } else { fs::write(&self.path, serialized)?; @@ -234,16 +242,16 @@ impl KvStore { { let key_str = key.to_string(); let serialized = serde_json::to_string(value)?; - + // Update in-memory data { let mut data = self.data.lock().unwrap(); data.insert(key_str, serialized); } - + // Save to disk self.save()?; - + Ok(()) } @@ -268,7 +276,7 @@ impl KvStore { Some(serialized) => { let value: V = serde_json::from_str(serialized)?; Ok(value) - }, + } None => Err(KvsError::KeyNotFound(key_str)), } } @@ -287,7 +295,7 @@ impl KvStore { K: ToString, { let key_str = key.to_string(); - + // Update in-memory data { let mut data = self.data.lock().unwrap(); @@ -295,10 +303,10 @@ impl KvStore { return Err(KvsError::KeyNotFound(key_str)); } } - + // Save to disk self.save()?; - + Ok(()) } @@ -317,7 +325,7 @@ impl KvStore { { let key_str = key.to_string(); let data = self.data.lock().unwrap(); - + Ok(data.contains_key(&key_str)) } @@ -328,7 +336,7 @@ impl KvStore { /// A vector of keys as strings pub fn keys(&self) -> Result> { let data = self.data.lock().unwrap(); - + Ok(data.keys().cloned().collect()) } @@ -343,10 +351,10 @@ impl KvStore { let mut data = self.data.lock().unwrap(); data.clear(); } - + // Save to disk self.save()?; - + Ok(()) }