Implement wallet encryption and decryption using AES-GCM; update save…

…/load functions to handle encrypted data

crates/command/src/keytool_cli/mod.rs

@@ -1,4 +1,4 @@
-use std::io;
+use std::io::{self, Write};
 use colored::Colorize;
 use key::{generate_karix_address, import_from_private_key, import_from_seed_phrase, list_wallet_files, load_wallet, save_wallet, set_selected_wallet  };
 
@@ -60,22 +60,49 @@ pub fn handle_keytool_command() -> Option<String> {
         let command = &args[2];
         // Use string comparison in the match statement
         match command.as_str() {
-            "generate" => { 
+            "generate" => {
                 println!("Enter mnemonic length (12 or 24):");
                 let mut mnemonic_length_str = String::new();
-                io::stdin().read_line(&mut mnemonic_length_str).unwrap();
-                let mnemonic_length: usize = mnemonic_length_str.trim().parse().expect("Invalid input");
-
-                let (private_key, public_address, seed_phrase) = generate_karix_address(mnemonic_length);
-                println!("New address generated:");
-                println!("Private Key: {}", private_key.green());
-                println!("Public Address: {}", public_address.green());
-                println!("Seed Phrase: {}", seed_phrase.green());
-
-                save_wallet(&public_address, &private_key, &seed_phrase);
-
-                return Some(public_address);
+                match io::stdin().read_line(&mut mnemonic_length_str) {
+                    Ok(_) => {
+                        match mnemonic_length_str.trim().parse::<usize>() {
+                            Ok(mnemonic_length) => {
+                                if mnemonic_length != 12 && mnemonic_length != 24 {
+                                    println!("{}", "Invalid mnemonic length. Must be 12 or 24.".red());
+                                    return None;
+                                }
+
+                                let (private_key, public_address, seed_phrase) = generate_karix_address(mnemonic_length);
+                                println!("New address generated:");
+                                println!("Private Key: {}", private_key.green());
+                                println!("Public Address: {}", public_address.green());
+                                println!("Seed Phrase: {}", seed_phrase.green());
+
+                                let password = prompt_password(true);
+                                match save_wallet(&public_address, &private_key, &seed_phrase, &password) {
+                                    Ok(_) => {
+                                        println!("Wallet saved successfully!");
+                                        return Some(public_address);
+                                    },
+                                    Err(e) => {
+                                        println!("{}", format!("Failed to save wallet: {}", e).red());
+                                        return None;
+                                    }
+                                }
+                            },
+                            Err(_) => {
+                                println!("{}", "Invalid input - please enter 12 or 24".red());
+                                return None;
+                            }
+                        }
+                    },
+                    Err(e) => {
+                        println!("{}", format!("Failed to read input: {}", e).red());
+                        return None;
+                    }
+                }
             },
+
             "balance" => {
                 println!("Enter public address:");
                 let mut public_address = String::new();
@@ -142,21 +169,32 @@ pub fn handle_keytool_command() -> Option<String> {
                 None
             },
 
-            "wallet" => { // String comparison for "wallet"
+            "wallet" => {
                 println!("Enter public address to load:");
                 let mut public_address = String::new();
-                io::stdin().read_line(&mut public_address).unwrap();
-                let public_address = public_address.trim().to_string();
-
-                if let Some(wallet_data) = load_wallet(&public_address) {
-                    println!("Wallet loaded:");
-                    println!("Address: {}", wallet_data.address.green());
-                    println!("Private Key: {}", wallet_data.private_key.green());
-                    println!("Seed Phrase: {}", wallet_data.seed_phrase.green());
-                    return Some(public_address);
-                } else {
-                    println!("Wallet not found for address: {}", public_address.red());
-                    return None; // Return None to indicate no address to be used further
+                match io::stdin().read_line(&mut public_address) {
+                    Ok(_) => {
+                        let public_address = public_address.trim().to_string();
+                        let password = prompt_password(false);
+
+                        match load_wallet(&public_address, &password) {
+                            Ok(wallet_data) => {
+                                println!("Wallet loaded:");
+                                println!("Address: {}", wallet_data.address.green());
+                                println!("Private Key: {}", wallet_data.private_key.green());
+                                println!("Seed Phrase: {}", wallet_data.seed_phrase.green());
+                                return Some(public_address);
+                            },
+                            Err(e) => {
+                                println!("{}", format!("Failed to load wallet: {}", e).red());
+                                return None;
+                            }
+                        }
+                    },
+                    Err(e) => {
+                        println!("{}", format!("Failed to read input: {}", e).red());
+                        return None;
+                    }
                 }
             },
 
@@ -232,13 +270,28 @@ pub fn handle_keytool_command() -> Option<String> {
 
                 match import_from_seed_phrase(phrase.trim()) {
                     Ok((private_key, _, public_address)) => {
-                        save_wallet(&public_address, &private_key, phrase.trim());
-                        set_selected_wallet(&public_address).expect("Failed to set selected wallet");
-                        println!("Imported wallet with address: {}", public_address);
-                        return Some(public_address);
+                        let password = prompt_password(true);
+                        match save_wallet(&public_address, &private_key, phrase.trim(), &password) {
+                            Ok(_) => {
+                                match set_selected_wallet(&public_address) {
+                                    Ok(_) => {
+                                        println!("Imported wallet with address: {}", public_address);
+                                        return Some(public_address);
+                                    },
+                                    Err(e) => {
+                                        println!("{}", format!("Failed to set selected wallet: {}", e).red());
+                                        return None;
+                                    }
+                                }
+                            },
+                            Err(e) => {
+                                println!("{}", format!("Failed to save wallet: {}", e).red());
+                                return None;
+                            }
+                        }
                     },
                     Err(e) => {
-                        println!("Failed to import seed phrase: {}", e);
+                        println!("{}", format!("Failed to import seed phrase: {}", e).red());
                         return None;
                     }
                 }
@@ -251,13 +304,28 @@ pub fn handle_keytool_command() -> Option<String> {
 
                 match import_from_private_key(private_key.trim()) {
                     Ok((private_key, _, public_address)) => {
-                        save_wallet(&public_address, &private_key, "");
-                        set_selected_wallet(&public_address).expect("Failed to set selected wallet");
-                        println!("Imported wallet with address: {}", public_address);
-                        return Some(public_address);
+                        let password = prompt_password(true);
+                        match save_wallet(&public_address, &private_key, "", &password) {
+                            Ok(_) => {
+                                match set_selected_wallet(&public_address) {
+                                    Ok(_) => {
+                                        println!("Imported wallet with address: {}", public_address);
+                                        return Some(public_address);
+                                    },
+                                    Err(e) => {
+                                        println!("{}", format!("Failed to set selected wallet: {}", e).red());
+                                        return None;
+                                    }
+                                }
+                            },
+                            Err(e) => {
+                                println!("{}", format!("Failed to save wallet: {}", e).red());
+                                return None;
+                            }
+                        }
                     },
                     Err(e) => {
-                        println!("Failed to import private key: {}", e);
+                        println!("{}", format!("Failed to import private key: {}", e).red());
                         return None;
                     }
                 }
@@ -272,3 +340,24 @@ pub fn handle_keytool_command() -> Option<String> {
         return None;
     }
 }
+
+
+fn prompt_password(confirm: bool) -> String {
+    print!("Enter password for wallet: ");
+    io::stdout().flush().unwrap();
+    let mut password = String::new();
+    io::stdin().read_line(&mut password).unwrap();
+    let password = password.trim().to_string();
+
+    if confirm {
+        print!("Confirm password: ");
+        io::stdout().flush().unwrap();
+        let mut confirm = String::new();
+        io::stdin().read_line(&mut confirm).unwrap();
+        if password != confirm.trim() {
+            println!("{}", "Passwords do not match!".red());
+            return prompt_password(true);
+        }
+    }
+    password
+}

crates/core/wallet/key/Cargo.toml

@@ -16,7 +16,7 @@ bip39.workspace = true
 secp256k1.workspace = true
 rand.workspace = true
 hex.workspace = true
-serde_json.workspace = true
+serde_json = { workspace = true , default-features = false, features = ["alloc"] }
 lazy_static.workspace = true
 consensus-pos.workspace = true
 
@@ -26,7 +26,8 @@ serde.workspace = true
 chacha20poly1305.workspace = true
 thiserror.workspace = true
 argon2.workspace = true
-
+aes-gcm = "0.10.3"
+base64 = "0.22.1"
 # This is a workspace crate, so we need to include the following dependencies
 k2.workspace = true
 move-core-types = { workspace = true }

crates/core/wallet/key/src/lib.rs

@@ -15,6 +15,46 @@ use thiserror::Error;
 use k2::{blockchain::get_kari_dir, config::{load_config, save_config}};
 use serde_json::json;
 
+use aes_gcm::{
+    aead::{Aead, KeyInit},
+    Aes256Gcm, Nonce,
+};
+use argon2::{
+    password_hash::{PasswordHasher, SaltString},
+    Argon2, PasswordHash,
+};
+use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
+
+
+#[derive(Error, Debug)]
+pub enum WalletError {
+    #[error("Encryption error: {0}")]
+    EncryptionError(String),
+    #[error("Decryption error: {0}")]
+    DecryptionError(String),
+    #[error("IO error: {0}")]
+    IoError(#[from] io::Error),
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct EncryptedData {
+    ciphertext: Vec<u8>,
+    salt: String,  
+    nonce: Vec<u8>,
+}
+
+fn derive_key(password: &str, salt: &SaltString) -> Result<[u8; 32], WalletError> {
+    let argon2 = Argon2::default();
+    let password_hash = argon2
+        .hash_password(password.as_bytes(), salt)
+        .map_err(|e| WalletError::EncryptionError(e.to_string()))?;
+    let mut key = [0u8; 32];
+    key.copy_from_slice(&password_hash.hash.unwrap().as_bytes()[0..32]);
+    Ok(key)
+}
+
+
+
 pub fn check_wallet_exists() -> bool {
     match list_wallet_files() {
         Ok(wallets) => !wallets.is_empty(),
@@ -29,24 +69,6 @@ pub struct Wallet {
     pub seed_phrase: String,
 }
 
-pub fn save_wallet(address: &str, private_key: &str, seed_phrase: &str) {
-    let kari_dir = get_kari_dir();
-    let wallet_dir = kari_dir.join("wallets");
-    fs::create_dir_all(&wallet_dir).expect("Unable to create wallets directory");
-
-    let wallet_file = wallet_dir.join(format!("{}.toml", address));
-    let wallet_data = Wallet {
-        address: address.to_string(),
-        private_key: private_key.to_string(),
-        seed_phrase: seed_phrase.to_string(),
-    };
-
-    let toml_string = toml::to_string(&wallet_data).expect("Unable to serialize wallet to TOML");
-    let mut file = fs::File::create(&wallet_file).expect("Unable to create wallet file");
-    file.write_all(toml_string.as_bytes()).expect("Unable to write wallet to file");
-    println!("Wallet saved to {:?}", wallet_file);
-}
-
 /// Set the selected wallet address in the configuration
 pub fn set_selected_wallet(wallet_address: &str) -> io::Result<()> {
     // Load existing config
@@ -61,36 +83,70 @@ pub fn set_selected_wallet(wallet_address: &str) -> io::Result<()> {
     save_config(&config)
 }
 
-pub fn load_wallet(address: &str) -> Option<Wallet> {
-    // Input validation with logging
-    if address.trim().is_empty() {
-        debug!("Attempted to load wallet with empty address");
-        return None;
-    }
+pub fn save_wallet(address: &str, private_key: &str, seed_phrase: &str, password: &str) -> Result<(), WalletError> {
+    let wallet_data = Wallet {
+        address: address.to_string(),
+        private_key: private_key.to_string(),
+        seed_phrase: seed_phrase.to_string(),
+    };
+
+    let salt = SaltString::generate(&mut OsRng);
+    let key = derive_key(password, &salt)?;
+    let cipher = Aes256Gcm::new_from_slice(&key).unwrap();
+    let binding = rand::random::<[u8; 12]>();
+    let nonce = Nonce::from_slice(&binding);
+
+    let toml_string = toml::to_string(&wallet_data)
+        .map_err(|e| WalletError::EncryptionError(e.to_string()))?;
+
+    let encrypted = cipher
+        .encrypt(nonce, toml_string.as_bytes())
+        .map_err(|e| WalletError::EncryptionError(e.to_string()))?;
+
+    let encrypted_data = EncryptedData {
+        ciphertext: encrypted,
+        salt: salt.to_string(),
+        nonce: nonce.to_vec(),
+    };
 
     let kari_dir = get_kari_dir();
-    let wallet_file: PathBuf = kari_dir.join("wallets").join(format!("{}.toml", address));
-
-    debug!("Attempting to load wallet from: {}", wallet_file.display());
-
-    if wallet_file.exists() {
-        // Handle potential IO and parsing errors with logging
-        fs::read_to_string(&wallet_file)
-            .map_err(|e| {
-                error!("Failed to read wallet file: {}", e);
-                e
-            })
-            .ok()
-            .and_then(|data| {
-                toml::from_str(&data).map_err(|e| {
-                    error!("Failed to parse wallet TOML: {}", e);
-                    e
-                }).ok()
-            })
-    } else {
-        debug!("Wallet file not found: {}", wallet_file.display());
-        None
-    }
+    let wallet_dir = kari_dir.join("wallets");
+    fs::create_dir_all(&wallet_dir)?;
+
+    let wallet_file = wallet_dir.join(format!("{}.enc", address));
+    let encrypted_json = serde_json::to_string(&encrypted_data)
+        .map_err(|e| WalletError::EncryptionError(e.to_string()))?;
+
+    fs::write(wallet_file, encrypted_json)?;
+    Ok(())
+}
+
+pub fn load_wallet(address: &str, password: &str) -> Result<Wallet, WalletError> {
+    let kari_dir = get_kari_dir();
+    let wallet_file = kari_dir.join("wallets").join(format!("{}.enc", address));
+
+    let encrypted_json = fs::read_to_string(wallet_file)?;
+    let encrypted_data: EncryptedData = serde_json::from_str(&encrypted_json)
+        .map_err(|e| WalletError::DecryptionError(e.to_string()))?;
+
+    let salt = SaltString::from_b64(&encrypted_data.salt)
+        .map_err(|e| WalletError::DecryptionError(e.to_string()))?;
+    let key = derive_key(password, &salt)?;
+
+    let cipher = Aes256Gcm::new_from_slice(&key).unwrap();
+    let nonce = Nonce::from_slice(&encrypted_data.nonce);
+
+    let decrypted = cipher
+        .decrypt(nonce, encrypted_data.ciphertext.as_slice())
+        .map_err(|e| WalletError::DecryptionError(e.to_string()))?;
+
+    let decrypted_str = String::from_utf8(decrypted)
+        .map_err(|e| WalletError::DecryptionError(e.to_string()))?;
+
+    let wallet: Wallet = toml::from_str(&decrypted_str)
+        .map_err(|e| WalletError::DecryptionError(e.to_string()))?;
+
+    Ok(wallet)
 }
 
 pub fn generate_karix_address(word_count: usize) -> (String, String, String) {