From bea9a26e3dcf6ed1dcc703848a942d343e38360c Mon Sep 17 00:00:00 2001
From: Aditya Kulkarni <adityapk@gmail.com>
Date: Tue, 22 Oct 2019 08:08:48 -0700
Subject: [PATCH] Recover seed properly

---
 lib/src/lightclient.rs | 64 ++++++++++++++++++++++++++++++++++++++++++
 src/main.rs            | 53 ++++++++++++----------------------
 2 files changed, 82 insertions(+), 35 deletions(-)

diff --git a/lib/src/lightclient.rs b/lib/src/lightclient.rs
index 1e570f0..0d17f8a 100644
--- a/lib/src/lightclient.rs
+++ b/lib/src/lightclient.rs
@@ -327,6 +327,47 @@ impl LightClient {
         Ok(lc)
     }
 
+    pub fn attempt_recover_seed(config: &LightClientConfig) -> Result<String, String> {
+        use std::io::prelude::*;
+        use byteorder::{LittleEndian, ReadBytesExt,};
+        use bip39::{Mnemonic, Language};
+        use zcash_primitives::serialize::Vector;
+
+        let mut reader = BufReader::new(File::open(config.get_wallet_path()).unwrap());
+        let version = reader.read_u64::<LittleEndian>().unwrap();
+        println!("Reading wallet version {}", version);
+
+        let encrypted = if version >= 4 {
+            reader.read_u8().unwrap() > 0
+        } else {
+            false
+        };
+
+        if encrypted {
+            return Err("The wallet is encrypted!".to_string());
+        }
+
+        let mut enc_seed = [0u8; 48];
+        if version >= 4 {
+            reader.read_exact(&mut enc_seed).unwrap();
+        }
+
+        let _nonce = if version >= 4 {
+            Vector::read(&mut reader, |r| r.read_u8()).unwrap()
+        } else {
+            vec![]
+        };
+
+        // Seed
+        let mut seed_bytes = [0u8; 32];
+        reader.read_exact(&mut seed_bytes).unwrap();
+
+        let phrase = Mnemonic::from_entropy(&seed_bytes, Language::English,).unwrap().phrase().to_string();
+
+        Ok(phrase)
+    }
+
+
     pub fn last_scanned_height(&self) -> u64 {
         self.wallet.read().unwrap().last_scanned_height() as u64
     }
@@ -1021,4 +1062,27 @@ pub mod tests {
         }
     }
 
+    #[test]
+    pub fn test_recover_seed() {
+        // Create a new tmp director
+        {
+            let tmp = TempDir::new("lctest").unwrap();
+            let dir_name = tmp.path().to_str().map(|s| s.to_string());
+
+            // A lightclient to a new, empty directory works.
+            let config = LightClientConfig::create_unconnected("test".to_string(), dir_name);
+            let lc = LightClient::new(&config, 0).unwrap();
+            let seed = lc.do_seed_phrase().unwrap()["seed"].as_str().unwrap().to_string();
+            lc.do_save().unwrap();
+
+            assert_eq!(seed, LightClient::attempt_recover_seed(&config).unwrap());
+
+            // Now encrypt and save the file
+            lc.wallet.write().unwrap().encrypt("password".to_string()).unwrap();
+            lc.do_save().unwrap();
+
+            assert!(LightClient::attempt_recover_seed(&config).is_err());
+        }
+    }
+
 }
\ No newline at end of file
diff --git a/src/main.rs b/src/main.rs
index 3222923..a9f96ae 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,4 @@
-use std::io::{Result, Error, ErrorKind};
+use std::io::{self, Error, ErrorKind};
 use std::sync::Arc;
 use std::sync::mpsc::{channel, Sender, Receiver};
 
@@ -20,7 +20,7 @@ use log4rs::append::rolling_file::policy::compound::{
 
 
 /// Build the Logging config
-fn get_log_config(config: &LightClientConfig) -> Result<Config> {
+fn get_log_config(config: &LightClientConfig) -> io::Result<Config> {
     let window_size = 3; // log0, log1, log2
     let fixed_window_roller =
         FixedWindowRoller::builder().build("zecwallet-light-wallet-log{}",window_size).unwrap();
@@ -91,7 +91,21 @@ pub fn main() {
                     .get_matches();
 
     if matches.is_present("recover") {
-        attempt_recover_seed();
+        // Create a Light Client Config in an attempt to recover the file.
+        let config = LightClientConfig {
+            server: "0.0.0.0:0".parse().unwrap(),
+            chain_name: "main".to_string(),
+            sapling_activation_height: 0,
+            consensus_branch_id: "000000".to_string(),
+            anchor_offset: 0,
+            no_cert_verification: false,
+            data_dir: None,
+        };
+
+        match LightClient::attempt_recover_seed(&config) {
+            Ok(seed) => println!("Recovered seed: '{}'", seed),
+            Err(e)   => eprintln!("Failed to recover seed. Error: {}", e)
+        };
         return;
     }
 
@@ -150,7 +164,7 @@ pub fn main() {
 }
 
 fn startup(server: http::Uri, dangerous: bool, seed: Option<String>, first_sync: bool, print_updates: bool)
-        -> Result<(Sender<(String, Vec<String>)>, Receiver<String>)> {
+        -> io::Result<(Sender<(String, Vec<String>)>, Receiver<String>)> {
     // Try to get the configuration
     let (config, latest_block_height) = LightClientConfig::create(server.clone(), dangerous)?;
 
@@ -301,34 +315,3 @@ fn command_loop(lightclient: Arc<LightClient>) -> (Sender<(String, Vec<String>)>
 
     (command_tx, resp_rx)
 }
-
-fn attempt_recover_seed() {
-    use std::fs::File;
-    use std::io::prelude::*;
-    use std::io::{BufReader};
-    use byteorder::{LittleEndian, ReadBytesExt,};
-    use bip39::{Mnemonic, Language};
-
-    // Create a Light Client Config in an attempt to recover the file.
-    let config = LightClientConfig {
-        server: "0.0.0.0:0".parse().unwrap(),
-        chain_name: "main".to_string(),
-        sapling_activation_height: 0,
-        consensus_branch_id: "000000".to_string(),
-        anchor_offset: 0,
-        no_cert_verification: false,
-        data_dir: None,
-    };
-
-    let mut reader = BufReader::new(File::open(config.get_wallet_path()).unwrap());
-    let version = reader.read_u64::<LittleEndian>().unwrap();
-    println!("Reading wallet version {}", version);
-
-    // Seed
-    let mut seed_bytes = [0u8; 32];
-    reader.read_exact(&mut seed_bytes).unwrap();
-
-    let phrase = Mnemonic::from_entropy(&seed_bytes, Language::English,).unwrap().phrase().to_string();
-
-    println!("Recovered seed phrase:\n{}", phrase);
-}