feat: use new argument to be able to set database path by commandline

replydev · Aug 30, 2024 · effd184 · effd184
1 parent 9fafde3
commit effd184
Show file tree

Hide file tree

Showing 5 changed files with 45 additions and 35 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -7,6 +7,7 @@ use interface::event::{Event, EventHandler};
 use interface::handlers::handle_key_events;
 use interface::ui::Tui;
 use otp::otp_element::{OTPDatabase, CURRENT_DATABASE_VERSION};
+use path::init_path;
 use ratatui::prelude::CrosstermBackend;
 use ratatui::Terminal;
 use reading::{get_elements_from_input, get_elements_from_stdin, ReadResult};
@@ -24,7 +25,9 @@ mod path;
 mod reading;
 mod utils;
 
-fn init(read_password_from_stdin: bool) -> color_eyre::Result<ReadResult> {
+fn init(args: &CotpArgs) -> color_eyre::Result<ReadResult> {
+    init_path(args);
+
     match utils::init_app() {
         Ok(first_run) => {
             if first_run {
@@ -38,7 +41,7 @@ fn init(read_password_from_stdin: bool) -> color_eyre::Result<ReadResult> {
                 let save_result = database.save_with_pw(&pw);
                 pw.zeroize();
                 save_result.map(|(key, salt)| (database, key, salt.to_vec()))
-            } else if read_password_from_stdin {
+            } else if args.password_from_stdin {
                 get_elements_from_stdin()
             } else {
                 get_elements_from_input()
@@ -52,7 +55,7 @@ fn main() -> AppResult<()> {
     color_eyre::install()?;
 
     let cotp_args: CotpArgs = CotpArgs::parse();
-    let (database, mut key, salt) = match init(cotp_args.password_from_stdin) {
+    let (database, mut key, salt) = match init(&cotp_args) {
         Ok(v) => v,
         Err(e) => {
             println!("{e}");

diff --git a/src/otp/otp_element.rs b/src/otp/otp_element.rs
@@ -4,7 +4,7 @@ use std::{fs::File, io::Write, vec};
 
 use crate::crypto::cryptography::{argon_derive_key, encrypt_string_with_key, gen_salt};
 use crate::otp::otp_error::OtpError;
-use crate::path::get_db_path;
+use crate::path::DATABASE_PATH;
 use data_encoding::BASE32_NOPAD;
 use qrcode::render::unicode;
 use qrcode::QrCode;
@@ -69,7 +69,7 @@ impl OTPDatabase {
     fn overwrite_database_key(&self, key: &Vec<u8>, salt: &[u8]) -> Result<(), std::io::Error> {
         let json: &str = &serde_json::to_string(&self)?;
         let encrypted = encrypt_string_with_key(json.to_string(), key, salt).unwrap();
-        let mut file = File::create(get_db_path())?;
+        let mut file = File::create(DATABASE_PATH.get().unwrap())?;
         match serde_json::to_string(&encrypted) {
             Ok(content) => {
                 file.write_all(content.as_bytes())?;

diff --git a/src/path.rs b/src/path.rs
@@ -3,16 +3,26 @@ use std::path::PathBuf;
 use std::sync::OnceLock;
 use std::{env, fs};
 
+use crate::arguments::CotpArgs;
+
 const CURRENT_DB_PATH: &str = "./db.cotp";
 const XDG_PATH: &str = "cotp/db.cotp";
 const HOME_PATH: &str = ".cotp/db.cotp";
 
-static ONCE_COMPUTED_PATH: OnceLock<PathBuf> = OnceLock::new();
+pub static DATABASE_PATH: OnceLock<PathBuf> = OnceLock::new();
 
-pub fn get_db_path() -> PathBuf {
-    env::var("COTP_DB_PATH")
-        .map(PathBuf::from)
-        .unwrap_or_else(|_| get_default_db_path())
+/// Initialize singleton database path
+pub fn init_path(args: &CotpArgs) -> PathBuf {
+    DATABASE_PATH
+        .get_or_init(|| {
+            args.database_path
+                .as_ref()
+                .map(String::from)
+                .or(env::var("COTP_DB_PATH").ok())
+                .map(PathBuf::from)
+                .unwrap_or_else(|| get_default_db_path())
+        })
+        .to_owned()
 }
 
 // Pushing an absolute path to a PathBuf replaces the entire PathBuf: https://doc.rust-lang.org/std/path/struct.PathBuf.html#method.push
@@ -26,26 +36,21 @@ fn get_default_db_path() -> PathBuf {
 
     let home_path = home_dir().map(|path| path.join(HOME_PATH));
 
-    ONCE_COMPUTED_PATH
-        .get_or_init(|| {
-            data_dir()
-                .map(PathBuf::from)
-                .map(|p| p.join(XDG_PATH))
-                .map(|xdg| {
-                    if !xdg.exists() {
-                        if let Some(home) = &home_path {
-                            if home.exists() {
-                                fs::create_dir_all(xdg.parent().unwrap())
-                                    .expect("Failed to create dir");
-                                fs::copy(home, xdg.as_path())
-                                    .expect("Failed on copy from legacy dir to XDG_DATA_HOME");
-                            }
-                        }
+    data_dir()
+        .map(PathBuf::from)
+        .map(|p| p.join(XDG_PATH))
+        .map(|xdg| {
+            if !xdg.exists() {
+                if let Some(home) = &home_path {
+                    if home.exists() {
+                        fs::create_dir_all(xdg.parent().unwrap()).expect("Failed to create dir");
+                        fs::copy(home, xdg.as_path())
+                            .expect("Failed on copy from legacy dir to XDG_DATA_HOME");
                     }
-                    xdg
-                })
-                .or(home_path)
-                .unwrap_or(portable_path)
+                }
+            }
+            xdg
         })
-        .to_owned()
+        .or(home_path)
+        .unwrap_or(portable_path)
 }
diff --git a/src/reading.rs b/src/reading.rs
@@ -1,6 +1,6 @@
 use crate::crypto;
 use crate::otp::otp_element::{OTPDatabase, OTPElement};
-use crate::path::get_db_path;
+use crate::path::DATABASE_PATH;
 use crate::utils;
 use color_eyre::eyre::{eyre, ErrReport};
 use std::fs::read_to_string;
@@ -28,7 +28,8 @@ fn get_elements_with_password(mut password: String) -> color_eyre::Result<ReadRe
 }
 
 pub fn read_decrypted_text(password: &str) -> color_eyre::Result<(String, Vec<u8>, Vec<u8>)> {
-    let encrypted_contents = read_to_string(get_db_path()).map_err(ErrReport::from)?;
+    let encrypted_contents =
+        read_to_string(DATABASE_PATH.get().unwrap()).map_err(ErrReport::from)?;
     if encrypted_contents.is_empty() {
         return match delete_db() {
             Ok(_) => Err(eyre!(
@@ -58,5 +59,5 @@ pub fn read_from_file(password: &str) -> color_eyre::Result<ReadResult> {
 }
 
 fn delete_db() -> io::Result<()> {
-    std::fs::remove_file(get_db_path())
+    std::fs::remove_file(DATABASE_PATH.get().unwrap())
 }
diff --git a/src/utils.rs b/src/utils.rs
@@ -1,8 +1,9 @@
-use crate::path::get_db_path;
 use std::time::{SystemTime, UNIX_EPOCH};
 
+use crate::path::DATABASE_PATH;
+
 pub fn init_app() -> Result<bool, ()> {
-    let db_path = get_db_path();
+    let db_path = DATABASE_PATH.get().unwrap(); // Safe to unwrap because we initialize
     let db_dir = db_path.parent().unwrap();
     if !db_dir.exists() {
         if let Err(_e) = std::fs::create_dir_all(db_dir) {