Wallet encryption commands

lib/src/commands.rs

@@ -201,6 +201,124 @@ impl Command for ExportCommand {
     }
 }
 
+struct EncryptCommand {}
+impl Command for EncryptCommand {
+    fn help(&self) -> String {
+        let mut h = vec![];
+        h.push("Encrypt the wallet with a password");
+        h.push("Note 1: This will encrypt the seed and the sapling and transparent private keys.");
+        h.push("        Use 'unlock' to temporarily unlock the wallet for spending or 'decrypt' ");
+        h.push("        to permanatly remove the encryption");
+        h.push("Note 2: If you forget the password, the only way to recover the wallet is to restore");
+        h.push("        from the seed phrase.");
+        h.push("Usage:");
+        h.push("encrypt password");
+        h.push("");
+        h.push("Example:");
+        h.push("encrypt my_strong_password");
+
+        h.join("\n")
+    }
+
+    fn short_help(&self) -> String {
+        "Encrypt the wallet with a password".to_string()
+    }
+
+    fn exec(&self, args: &[&str], lightclient: &LightClient) -> String {
+        if args.len() != 1 {
+            return self.help();
+        }
+
+        let passwd = args[0].to_string();
+
+        match lightclient.wallet.write().unwrap().encrypt(passwd) {
+            Ok(_)  => object!{ "result" => "success" },
+            Err(e) => object!{
+                "result" => "error",
+                "error"  => e.to_string()
+            }
+        }.pretty(2)
+    }
+}
+
+struct DecryptCommand {}
+impl Command for DecryptCommand {
+    fn help(&self) -> String {
+        let mut h = vec![];
+        h.push("Completely remove wallet encryption, storing the wallet in plaintext on disk");
+        h.push("Note 1: This will decrypt the seed and the sapling and transparent private keys and store them on disk.");
+        h.push("        Use 'unlock' to temporarily unlock the wallet for spending");
+        h.push("Note 2: If you've forgotten the password, the only way to recover the wallet is to restore");
+        h.push("        from the seed phrase.");
+        h.push("Usage:");
+        h.push("decrypt password");
+        h.push("");
+        h.push("Example:");
+        h.push("decrypt my_strong_password");
+
+        h.join("\n")
+    }
+
+    fn short_help(&self) -> String {
+        "Completely remove wallet encryption".to_string()
+    }
+
+    fn exec(&self, args: &[&str], lightclient: &LightClient) -> String {
+        if args.len() != 1 {
+            return self.help();
+        }
+
+        let passwd = args[0].to_string();
+
+        match lightclient.wallet.write().unwrap().remove_encryption(passwd) {
+            Ok(_)  => object!{ "result" => "success" },
+            Err(e) => object!{
+                "result" => "error",
+                "error"  => e.to_string()
+            }
+        }.pretty(2)
+    }
+}
+
+
+struct UnlockCommand {}
+impl Command for UnlockCommand {
+    fn help(&self) -> String {
+        let mut h = vec![];
+        h.push("Unlock the wallet's encryption in memory, allowing spending from this wallet.");
+        h.push("Note 1: This will decrypt spending keys in memory only. The wallet remains encrypted on disk");
+        h.push("        Use 'decrypt' to remove the encryption permanatly.");
+        h.push("Note 2: If you've forgotten the password, the only way to recover the wallet is to restore");
+        h.push("        from the seed phrase.");
+        h.push("Usage:");
+        h.push("unlock password");
+        h.push("");
+        h.push("Example:");
+        h.push("unlock my_strong_password");
+
+        h.join("\n")
+    }
+
+    fn short_help(&self) -> String {
+        "Unlock wallet encryption for spending".to_string()
+    }
+
+    fn exec(&self, args: &[&str], lightclient: &LightClient) -> String {
+        if args.len() != 1 {
+            return self.help();
+        }
+
+        let passwd = args[0].to_string();
+
+        match lightclient.wallet.write().unwrap().unlock(passwd) {
+            Ok(_)  => object!{ "result" => "success" },
+            Err(e) => object!{
+                "result" => "error",
+                "error"  => e.to_string()
+            }
+        }.pretty(2)
+    }
+}
 
 struct SendCommand {}
 impl Command for SendCommand {
@@ -527,6 +645,9 @@ pub fn get_commands() -> Box<HashMap<String, Box<dyn Command>>> {
     map.insert("notes".to_string(),         Box::new(NotesCommand{}));
     map.insert("new".to_string(),           Box::new(NewAddressCommand{}));
     map.insert("seed".to_string(),          Box::new(SeedCommand{}));
+    map.insert("encrypt".to_string(),       Box::new(EncryptCommand{}));
+    map.insert("decrypt".to_string(),       Box::new(DecryptCommand{}));
+    map.insert("unlock".to_string(),        Box::new(UnlockCommand{}));
     map.insert("fixbip39bug".to_string(),   Box::new(FixBip39BugCommand{}));
 
     Box::new(map)

lib/src/lightclient.rs

@@ -248,6 +248,13 @@ impl LightClient {
 
     // Export private keys
     pub fn do_export(&self, addr: Option<String>) -> JsonValue {
+        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
+            error!("Wallet is locked");
+            return object!{
+                "error" => "Wallet is locked"
+            };
+        }
+
         // Clone address so it can be moved into the closure
         let address = addr.clone();
         let wallet = self.wallet.read().unwrap();
@@ -332,6 +339,22 @@ impl LightClient {
     }
 
     pub fn do_save(&self) -> Result<(), String> {
+        
+        // If the wallet is encrypted but unlocked, lock it again.
+        {
+            let mut wallet = self.wallet.write().unwrap();
+            if wallet.is_encrypted() && wallet.is_unlocked_for_spending() {
+                match wallet.lock() {
+                    Ok(_) => {},
+                    Err(e) => {
+                        let err = format!("ERR: {}", e);
+                        error!("{}", err);
+                        return Err(e.to_string());
+                    }
+                }
+            }
+        }        
+
         let mut file_buffer = BufWriter::with_capacity(
             1_000_000, // 1 MB write buffer
             File::create(self.config.get_wallet_path()).unwrap());
@@ -369,6 +392,13 @@ impl LightClient {
     }
 
     pub fn do_seed_phrase(&self) -> JsonValue {
+        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
+            error!("Wallet is locked");
+            return object!{
+                "error" => "Wallet is locked"
+            };
+        }
+
         let wallet = self.wallet.read().unwrap();
         object!{
             "seed"     => wallet.get_seed_phrase(),
@@ -549,6 +579,13 @@ impl LightClient {
 
     /// Create a new address, deriving it from the seed.
     pub fn do_new_address(&self, addr_type: &str) -> JsonValue {
+        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
+            error!("Wallet is locked");
+            return object!{
+                "error" => "Wallet is locked"
+            };
+        }
+
         let wallet = self.wallet.write().unwrap();
 
         let new_address = match addr_type {
@@ -784,6 +821,11 @@ impl LightClient {
     }
 
     pub fn do_send(&self, addrs: Vec<(&str, u64, Option<String>)>) -> Result<String, String> {
+        if !self.wallet.read().unwrap().is_unlocked_for_spending() {
+            error!("Wallet is locked");
+            return Err("Wallet is locked".to_string());
+        }
+
         info!("Creating transaction");
 
         let rawtx = self.wallet.write().unwrap().send_to_address(

lib/src/lightwallet.rs

@@ -91,7 +91,14 @@ impl ToBase58Check for [u8] {
 }
 
 pub struct LightWallet {
-    locked: bool,       // Is the wallet's spending keys locked?
+    // Is the wallet encrypted? If it is, then when writing to disk, the seed is always encrypted 
+    // and the individual spending keys are not written    
+    encrypted: bool,       
+
+    // In memory only (i.e, this field is not written to disk). Is the wallet unlocked and are
+    // the spending keys present to allow spending from this wallet?
+    unlocked: bool,
+
     enc_seed: [u8; 48], // If locked, this contains the encrypted seed
     nonce: Vec<u8>,     // Nonce used to encrypt the wallet. 
 
@@ -184,7 +191,8 @@ impl LightWallet {
             = LightWallet::get_zaddr_from_bip39seed(&config, &bip39_seed.as_bytes(), 0);
 
         Ok(LightWallet {
-            locked:     false,
+            encrypted:  false,
+            unlocked:   true,
             enc_seed:   [0u8; 48],
             nonce:      vec![],
             seed:       seed_bytes,
@@ -210,7 +218,7 @@ impl LightWallet {
 
         info!("Reading wallet version {}", version);
 
-        let locked = if version >= 4 {
+        let encrypted = if version >= 4 {
             reader.read_u8()? > 0
         } else {
             false
@@ -281,7 +289,8 @@ impl LightWallet {
         let birthday = reader.read_u64::<LittleEndian>()?;
 
         Ok(LightWallet{
-            locked:     locked,
+            encrypted:  encrypted,
+            unlocked:   !encrypted, // When reading from disk, if wallet is encrypted, it starts off locked. 
             enc_seed:   enc_seed,
             nonce:      nonce,
             seed:       seed_bytes,
@@ -298,11 +307,16 @@ impl LightWallet {
     }
 
     pub fn write<W: Write>(&self, mut writer: W) -> io::Result<()> {
+        if self.encrypted && self.unlocked {
+            return Err(Error::new(ErrorKind::InvalidInput, 
+                        format!("Cannot write while wallet is unlocked while encrypted.")));
+        }
+
         // Write the version
         writer.write_u64::<LittleEndian>(LightWallet::serialized_version())?;
 
         // Write if it is locked
-        writer.write_u8(if self.locked {1} else {0})?;
+        writer.write_u8(if self.encrypted {1} else {0})?;
 
         // Write the encrypted seed bytes
         writer.write_all(&self.enc_seed)?;
@@ -400,6 +414,10 @@ impl LightWallet {
     /// at the next position and add it to the wallet.
     /// NOTE: This does NOT rescan
     pub fn add_zaddr(&self) -> String {
+        if !self.unlocked {
+            return "".to_string();
+        }
+
         let pos = self.extsks.read().unwrap().len() as u32;
         let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&self.seed, Language::English).unwrap(), "");
 
@@ -418,6 +436,10 @@ impl LightWallet {
     /// at the next position.
     /// NOTE: This is not rescan the wallet
     pub fn add_taddr(&self) -> String {
+        if !self.unlocked {
+            return "".to_string();
+        }
+
         let pos = self.tkeys.read().unwrap().len() as u32;
         let bip39_seed = bip39::Seed::new(&Mnemonic::from_entropy(&self.seed, Language::English).unwrap(), "");
         
@@ -562,16 +584,20 @@ impl LightWallet {
     }
 
     pub fn get_seed_phrase(&self) -> String {
+        if !self.unlocked {
+            return "".to_string();
+        }
+
         Mnemonic::from_entropy(&self.seed, 
                                 Language::English,
         ).unwrap().phrase().to_string()
     }
 
-    pub fn lock(&mut self, passwd: String) -> io::Result<()> {
+    pub fn encrypt(&mut self, passwd: String) -> io::Result<()> {
         use sodiumoxide::crypto::secretbox;
 
-        if self.locked {
-            return Err(io::Error::new(ErrorKind::AlreadyExists, "Wallet is already locked"));
+        if self.encrypted && !self.unlocked {
+            return Err(io::Error::new(ErrorKind::AlreadyExists, "Wallet is already encrypted and locked"));
         }
 
         // Get the doublesha256 of the password, which is the right length
@@ -584,20 +610,32 @@ impl LightWallet {
         self.nonce = vec![];
         self.nonce.extend_from_slice(nonce.as_ref());
 
+        self.encrypted = true;
+        self.lock()?;
+
+        Ok(())
+    }
+
+    pub fn lock(&mut self) -> io::Result<()> {
         // Empty the seed and the secret keys
         self.seed.copy_from_slice(&[0u8; 32]);
         self.tkeys = Arc::new(RwLock::new(vec![]));
         self.extsks = Arc::new(RwLock::new(vec![]));
 
-        self.locked = true;
+        self.unlocked = false;
+
         Ok(())
     }
 
     pub fn unlock(&mut self, passwd: String) -> io::Result<()> {
-         use sodiumoxide::crypto::secretbox;
+        use sodiumoxide::crypto::secretbox;
 
-        if !self.locked {
-            return Err(io::Error::new(ErrorKind::AlreadyExists, "Wallet is not locked"));
+        if !self.encrypted {
+            return Err(Error::new(ErrorKind::AlreadyExists, "Wallet is not encrypted"));
+        }
+
+        if self.encrypted && self.unlocked {
+            return Err(Error::new(ErrorKind::AlreadyExists, "Wallet is already unlocked"));
         }
 
         // Get the doublesha256 of the password, which is the right length
@@ -650,18 +688,45 @@ impl LightWallet {
             tkeys.push(sk);
         }
 
-        // Everything checks out, so we'll update our wallet with the unlocked values
+        // Everything checks out, so we'll update our wallet with the decrypted values
         self.extsks = Arc::new(RwLock::new(extsks));
         self.tkeys = Arc::new(RwLock::new(tkeys));
         self.seed.copy_from_slice(&seed);
-        
-        self.nonce = vec![];
-        self.enc_seed.copy_from_slice(&[0u8; 48]);
-        self.locked = false;
+                
+        self.encrypted = true;
+        self.unlocked = true;
 
         Ok(())
     }
 
+    // Removing encryption means unlocking it and setting the self.encrypted = false,
+    // permanantly removing the encryption
+    pub fn remove_encryption(&mut self, passwd: String) -> io::Result<()> {        
+        if !self.encrypted {
+            return Err(Error::new(ErrorKind::AlreadyExists, "Wallet is not encrypted"));
+        }
+
+        // Unlock the wallet if it's locked
+        if !self.unlocked {
+            self.unlock(passwd)?;
+        }
+        
+        // Permanantly remove the encryption
+        self.encrypted = false;
+        self.nonce = vec![];
+        self.enc_seed.copy_from_slice(&[0u8; 48]);
+
+        Ok(())
+    }
+
+    pub fn is_encrypted(&self) -> bool {
+        return self.encrypted;
+    }
+
+    pub fn is_unlocked_for_spending(&self) -> bool {
+        return self.unlocked;
+    }
+
     pub fn zbalance(&self, addr: Option<String>) -> u64 {
         self.txs.read().unwrap()
             .values()
@@ -1245,6 +1310,10 @@ impl LightWallet {
         output_params: &[u8],
         tos: Vec<(&str, u64, Option<String>)>
     ) -> Result<Box<[u8]>, String> {
+        if !self.unlocked {
+            return Err("Cannot spend while wallet is locked".to_string());
+        }
+
         let start_time = now();
 
         let total_value = tos.iter().map(|to| to.1).sum::<u64>();
@@ -3083,10 +3152,10 @@ pub mod tests {
 
         let seed = wallet.seed;
 
-        wallet.lock("somepassword".to_string()).unwrap();
+        wallet.encrypt("somepassword".to_string()).unwrap();
 
-        // Locking a locked wallet should fail
-        assert!(wallet.lock("somepassword".to_string()).is_err());
+        // Encrypting an already encrypted wallet should fail
+        assert!(wallet.encrypt("somepassword".to_string()).is_err());
 
         // Serialize a locked wallet
         let mut serialized_data = vec![];
@@ -3120,6 +3189,12 @@ pub mod tests {
         // Unlocking an already unlocked wallet should fail
         assert!(wallet.unlock("somepassword".to_string()).is_err());
 
+        // Trying to serialize a encrypted but unlocked wallet should fail
+        assert!(wallet.write(&mut vec![]).is_err());
+
+        // ...but if we lock it again, it should serialize
+        wallet.lock().unwrap();
+        wallet.write(&mut vec![]).expect("Serialize wallet");
 
         // Try from a deserialized, locked wallet
         let mut wallet2 = LightWallet::read(&serialized_data[..], &config).unwrap();