From a6483cf333e6a5f3a0d48317b50c6304cfd956bb Mon Sep 17 00:00:00 2001
From: Awiteb <a@4rs.nl>
Date: Wed, 8 May 2024 05:02:04 +0300
Subject: [PATCH] feat: Support export and import with different password

---
 src/cli/export_command.rs | 31 +++++++++++++++++++++++++++----
 src/cli/import_command.rs | 29 +++++++++++++++++++++++++----
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/src/cli/export_command.rs b/src/cli/export_command.rs
index d488a97..78c436d 100644
--- a/src/cli/export_command.rs
+++ b/src/cli/export_command.rs
@@ -17,8 +17,10 @@
 use std::{fs, io::Error as IoError, io::ErrorKind as IoErrorKind, path::PathBuf};
 
 use clap::Args;
+use sha2::Digest;
 
 use crate::{
+    utils,
     vault::{BitWardenPasswords, Format, Vaults},
     LprsCommand,
     LprsError,
@@ -30,12 +32,17 @@ use crate::{
 /// Export command, used to export the vaults in `lprs` format or `BitWarden`
 /// format. The exported file will be a json file.
 pub struct Export {
+    // TODO: `force` flag to write on existing file
     /// The path to export to
-    path:   PathBuf,
+    path:                PathBuf,
     /// Format to export vaults in
     #[arg(short, long, value_name = "FORMAT", default_value_t= Format::Lprs)]
-    format: Format,
-    // TODO: `force` flag to write on existing file
+    format:              Format,
+    /// Encryption password of the exported vaults (in `lprs` format)
+    /// if there is not, will use the master password
+    #[arg(short = 'p', long)]
+    #[allow(clippy::option_option)]
+    encryption_password: Option<Option<String>>,
 }
 
 impl LprsCommand for Export {
@@ -46,8 +53,19 @@ impl LprsCommand for Export {
             self.path.display(),
             self.format
         );
+
+        let encryption_key: Option<[u8; 32]> =
+            utils::user_password(self.encryption_password, "Encryption Password:")?
+                .map(|p| sha2::Sha256::digest(p).into());
+
         let exported_data = match self.format {
-            Format::Lprs => vault_manager.json_export()?,
+            Format::Lprs => {
+                vault_manager.json_export(
+                    encryption_key
+                        .as_ref()
+                        .unwrap_or(&vault_manager.master_password),
+                )?
+            }
             Format::BitWarden => serde_json::to_string(&BitWardenPasswords::from(vault_manager))?,
         };
 
@@ -77,6 +95,11 @@ impl LprsCommand for Export {
                 format!("file `{}` is a directory", self.path.display()),
             )));
         }
+        if self.encryption_password.is_some() && self.format != Format::Lprs {
+            return Err(LprsError::Other(
+                "You only can to use the encryption password with `lprs` format".to_owned(),
+            ));
+        }
 
         Ok(())
     }
diff --git a/src/cli/import_command.rs b/src/cli/import_command.rs
index 700a8a1..977b677 100644
--- a/src/cli/import_command.rs
+++ b/src/cli/import_command.rs
@@ -22,8 +22,10 @@ use std::{
 };
 
 use clap::Args;
+use sha2::Digest;
 
 use crate::{
+    utils,
     vault::{BitWardenPasswords, Format, Vault, Vaults},
     LprsCommand,
     LprsError,
@@ -40,7 +42,12 @@ pub struct Import {
 
     /// The format to import from
     #[arg(short, long, default_value_t = Format::Lprs)]
-    format: Format,
+    format:              Format,
+    /// Decryption password of the imported vaults (in `lprs` format)
+    /// if there is not, will use the master password
+    #[arg(short = 'p', long)]
+    #[allow(clippy::option_option)]
+    decryption_password: Option<Option<String>>,
 }
 
 impl LprsCommand for Import {
@@ -52,10 +59,18 @@ impl LprsCommand for Import {
             vault_manager.vaults_file.display()
         );
 
+        let decryption_key: Option<[u8; 32]> =
+            utils::user_password(self.decryption_password, "Decryption password:")?
+                .map(|p| sha2::Sha256::digest(p).into());
+
         let imported_passwords_len = match self.format {
             Format::Lprs => {
-                let vaults =
-                    Vaults::json_reload(&vault_manager.master_password, &fs::read(self.path)?)?;
+                let vaults = Vaults::json_reload(
+                    decryption_key
+                        .as_ref()
+                        .unwrap_or(&vault_manager.master_password),
+                    &fs::read(self.path)?,
+                )?;
                 let vaults_len = vaults.len();
 
                 vault_manager.vaults.extend(vaults);
@@ -81,7 +96,7 @@ impl LprsCommand for Import {
     }
 
     fn validate_args(&self) -> LprsResult<()> {
-        if self
+        if !self
             .path
             .extension()
             .is_some_and(|e| e.to_string_lossy().eq_ignore_ascii_case("json"))
@@ -103,6 +118,12 @@ impl LprsCommand for Import {
                 format!("file `{}` is a directory", self.path.display()),
             )));
         }
+        if self.decryption_password.is_some() && self.format != Format::Lprs {
+            return Err(LprsError::Other(
+                "You only can to use the decryption password with `lprs` format".to_owned(),
+            ));
+        }
+
 
         Ok(())
     }