Merge pull request #183 from replydev/fix/import_on_first_run

Fix panic when importing on first run.

Author: replydev

src/argument_functions.rs

@@ -152,7 +152,7 @@ pub fn export(matches: &ArgMatches, database: &mut OTPDatabase) -> Result<String
 pub fn change_password(database: &mut OTPDatabase) -> Result<String, String> {
     let mut new_password = utils::verified_password("New password: ", 8);
     let r = match database.save_with_pw(&new_password) {
-        Ok(()) => Ok(String::from("Password changed")),
+        Ok(_) => Ok(String::from("Password changed")),
         Err(e) => Err(format!("An error has occurred: {e}")),
     };
     new_password.zeroize();

src/main.rs

@@ -31,13 +31,14 @@ fn init() -> Result<ReadResult, String> {
                     elements: vec![],
                     ..Default::default()
                 };
-                if database.save_with_pw(&pw).is_err() {
-                    return Err(String::from(
+                let save_result = database.save_with_pw(&pw);
+                pw.zeroize();
+                match save_result {
+                    Ok((key, salt)) => Ok((database, key, salt.to_vec())),
+                    Err(_) => Err(String::from(
                         "An error occurred during database overwriting",
-                    ));
+                    )),
                 }
-                pw.zeroize();
-                Ok((database, vec![], vec![]))
             } else {
                 get_elements()
             }

src/otp/otp_element.rs

@@ -70,10 +70,11 @@ impl OTPDatabase {
         }
     }
 
-    pub fn save_with_pw(&mut self, password: &str) -> Result<(), String> {
+    pub fn save_with_pw(&mut self, password: &str) -> Result<(Vec<u8>, [u8; 16]), String> {
         let salt = gen_salt()?;
         let key = argon_derive_key(password.as_bytes(), &salt)?;
-        self.save(&key, &salt)
+        self.save(&key, &salt)?;
+        Ok((key, salt))
     }
 
     pub fn export(&self, path: PathBuf) -> Result<PathBuf, String> {