feat: improve decrypt

src/cmd_decrypt.rs

@@ -17,7 +17,7 @@ use yubikey::piv::{AlgorithmId, decrypt_data};
 use yubikey::YubiKey;
 use zeroize::Zeroize;
 
-use crate::{cmd_encrypt, consts, crypto_simple, util, util_enc_file, util_envelop, util_file, util_pgp, util_piv};
+use crate::{cmd_encrypt, consts, crypto_simple, util, util_enc_file, util_env, util_envelop, util_file, util_pgp, util_piv};
 use crate::compress::GzStreamDecoder;
 use crate::config::TinyEncryptConfig;
 use crate::consts::{
@@ -40,6 +40,9 @@ pub struct CmdDecrypt {
     /// PIN
     #[arg(long, short = 'p')]
     pub pin: Option<String>,
+    /// KeyID
+    #[arg(long, short = 'k')]
+    pub key_id: Option<String>,
     /// Slot
     #[arg(long, short = 's')]
     pub slot: Option<String>,
@@ -81,18 +84,32 @@ pub fn decrypt(cmd_decrypt: CmdDecrypt) -> XResult<()> {
     let mut succeed_count = 0;
     let mut failed_count = 0;
     let mut total_len = 0_u64;
+    if cmd_decrypt.edit_file && (cmd_decrypt.paths.len() != 1) {
+        return simple_error!("Edit mode only allows one file assigned.");
+    }
+    let pin = match &cmd_decrypt.pin {
+        Some(pin) => Some(pin.clone()),
+        None => util_env::get_pin(),
+    };
+    let key_id = match &cmd_decrypt.key_id {
+        Some(key_id) => Some(key_id.clone()),
+        None => util_env::get_key_id(),
+    };
+
     for path in &cmd_decrypt.paths {
         let start_decrypt_single = Instant::now();
-        match decrypt_single(&config, path, &cmd_decrypt.pin, &cmd_decrypt.slot, &cmd_decrypt) {
+        match decrypt_single(&config, path, &pin, &key_id, &cmd_decrypt.slot, &cmd_decrypt) {
             Ok(len) => {
                 succeed_count += 1;
-                total_len += len;
-                success!(
-                    "Decrypt {} succeed, cost {} ms, file size {}",
-                    path.to_str().unwrap_or("N/A"),
-                    start_decrypt_single.elapsed().as_millis(),
-                    util_size::get_display_size(len as i64)
-                );
+                if len > 0 {
+                    total_len += len;
+                    success!(
+                        "Decrypt {} succeed, cost {} ms, file size {}",
+                        path.to_str().unwrap_or("N/A"),
+                        start_decrypt_single.elapsed().as_millis(),
+                        util_size::get_display_size(len as i64)
+                    );
+                }
             }
             Err(e) => {
                 failed_count += 1;
@@ -115,13 +132,14 @@ pub fn decrypt(cmd_decrypt: CmdDecrypt) -> XResult<()> {
 pub fn decrypt_single(config: &Option<TinyEncryptConfig>,
                       path: &PathBuf,
                       pin: &Option<String>,
+                      key_id: &Option<String>,
                       slot: &Option<String>,
                       cmd_decrypt: &CmdDecrypt) -> XResult<u64> {
     let path_display = format!("{}", path.display());
     util::require_tiny_enc_file_and_exists(path)?;
 
     let mut file_in = opt_result!(File::open(path), "Open file: {} failed: {}", &path_display);
-    let (_, meta) = opt_result!(
+    let (_, is_meta_compressed, meta) = opt_result!(
         util_enc_file::read_tiny_encrypt_meta_and_normalize(&mut file_in), "Read file: {}, failed: {}", &path_display);
     util_msg::when_debug(|| {
         debugging!("Found meta: {}", serde_json::to_string_pretty(&meta).unwrap());
@@ -139,7 +157,7 @@ pub fn decrypt_single(config: &Option<TinyEncryptConfig>,
     let digest_algorithm = cmd_decrypt.digest_algorithm.as_deref().unwrap_or("sha256");
     if cmd_decrypt.digest_file { DigestWrite::from_algo(digest_algorithm)?; } // FAST CHECK
 
-    let selected_envelop = select_envelop(&meta, config)?;
+    let selected_envelop = select_envelop(&meta, key_id, config)?;
 
     let key = SecVec(try_decrypt_key(config, selected_envelop, pin, slot)?);
     let nonce = SecVec(opt_result!(util::decode_base64(&meta.nonce), "Decode nonce failed: {}"));
@@ -166,45 +184,26 @@ pub fn decrypt_single(config: &Option<TinyEncryptConfig>,
 
     // Edit file
     if cmd_decrypt.edit_file {
-        let output = match decrypt_limited_content_to_vec(&mut file_in, &meta, cryptor, &key_nonce)? {
+        let file_content = match decrypt_limited_content_to_vec(&mut file_in, &meta, cryptor, &key_nonce)? {
             None => return Ok(0),
             Some(output) => output,
         };
-        let editor = match env::var("EDITOR") {
-            Ok(editor) => editor,
-            Err(_) => return simple_error!("EDITOR not found."),
-        };
-
-        // write file to temp file
-        let temp_dir = temp_dir();
-        let current_millis = util_time::get_current_millis();
-        let temp_file = temp_dir.join(&format!("tmp_file_{}_{}", current_millis, path_out));
-        information!("Temp file: {}", temp_file.display());
-        opt_result!(fs::write(&temp_file, &output), "Write temp file failed: {}");
+        let editor = get_file_editor();
+        let temp_file = create_edit_temp_file(&file_content, path_out)?;
 
         let do_edit_file = || -> XResult<()> {
-            // run editor
-            let mut command = Command::new(&editor);
-            command.arg(temp_file.to_str().expect("Get temp file path failed."));
-            let run_cmd_result = util_cmd::run_command_and_wait(&mut command);
-            debugging!("Run cmd result: {:?}", run_cmd_result);
-            let run_cmd_exit_status = opt_result!(run_cmd_result, "Run cmd {} failed: {}", editor);
-            if !run_cmd_exit_status.success() {
-                return simple_error!("Run cmd {} failed: {:?}", editor, run_cmd_exit_status.code());
-            }
-            let temp_file_content = opt_result!(fs::read_to_string(&temp_file), "Read file failed: {}");
-            if temp_file_content == output {
+            let temp_file_content = run_file_editor_and_wait_content(&editor, &temp_file)?;
+            if temp_file_content == file_content {
                 information!("Temp file is not changed.");
                 return Ok(());
             }
             success!("Temp file is changed, save file ...");
-
             drop(file_in);
             let mut meta = meta;
             meta.file_length = temp_file_content.len() as u64;
             meta.file_last_modified = util_time::get_current_millis() as u64;
             let mut file_out = File::create(path)?;
-            let _ = util_enc_file::write_tiny_encrypt_meta(&mut file_out, &meta, true)?;
+            let _ = util_enc_file::write_tiny_encrypt_meta(&mut file_out, &meta, is_meta_compressed)?;
             let compress_level = iff!(meta.compress, Some(Compression::default().level()), None);
             cmd_encrypt::encrypt_file(
                 &mut temp_file_content.as_bytes(), meta.file_length, &mut file_out, cryptor,
@@ -261,6 +260,37 @@ pub fn decrypt_single(config: &Option<TinyEncryptConfig>,
     Ok(meta.file_length)
 }
 
+fn run_file_editor_and_wait_content(editor: &str, temp_file: &PathBuf) -> XResult<String> {
+    let mut command = Command::new(editor);
+    command.arg(temp_file.to_str().expect("Get temp file path failed."));
+    let run_cmd_result = util_cmd::run_command_and_wait(&mut command);
+    debugging!("Run cmd result: {:?}", run_cmd_result);
+    let run_cmd_exit_status = opt_result!(run_cmd_result, "Run cmd {} failed: {}", editor);
+    if !run_cmd_exit_status.success() {
+        return simple_error!("Run cmd {} failed: {:?}", editor, run_cmd_exit_status.code());
+    }
+    Ok(opt_result!(fs::read_to_string(&temp_file), "Read file failed: {}"))
+}
+
+fn get_file_editor() -> String {
+    match env::var("EDITOR") {
+        Ok(editor) => editor,
+        Err(_) => {
+            warning!("EDITOR is not assigned, use default editor vi");
+            "vi".to_string()
+        }
+    }
+}
+
+fn create_edit_temp_file(file_content: &str, path_out: &str) -> XResult<PathBuf> {
+    let temp_dir = temp_dir();
+    let current_millis = util_time::get_current_millis();
+    let temp_file = temp_dir.join(&format!("tmp_file_{}_{}", current_millis, path_out));
+    information!("Temp file: {}", temp_file.display());
+    opt_result!(fs::write(&temp_file, file_content), "Write temp file failed: {}");
+    Ok(temp_file)
+}
+
 fn decrypt_limited_content_to_vec(mut file_in: &mut File,
                                   meta: &TinyEncryptMeta, cryptor: Cryptor, key_nonce: &KeyNonce) -> XResult<Option<String>> {
     if meta.file_length > 100 * 1024 {
@@ -450,7 +480,7 @@ fn try_decrypt_key_pgp(envelop: &TinyEncryptEnvelop, pin: &Option<String>) -> XR
     Ok(key)
 }
 
-fn select_envelop<'a>(meta: &'a TinyEncryptMeta, config: &Option<TinyEncryptConfig>) -> XResult<&'a TinyEncryptEnvelop> {
+fn select_envelop<'a>(meta: &'a TinyEncryptMeta, key_id: &Option<String>, config: &Option<TinyEncryptConfig>) -> XResult<&'a TinyEncryptEnvelop> {
     let envelops = match &meta.envelops {
         None => return simple_error!("No envelops found"),
         Some(envelops) => if envelops.is_empty() {
@@ -468,6 +498,25 @@ fn select_envelop<'a>(meta: &'a TinyEncryptMeta, config: &Option<TinyEncryptConf
         return Ok(selected_envelop);
     }
 
+    if let Some(key_id) = key_id {
+        for envelop in envelops {
+            let is_sid_matched = match config {
+                None => false,
+                Some(config) => match config.find_by_kid(&envelop.kid) {
+                    None => false,
+                    Some(config) => match &config.sid {
+                        None => false,
+                        Some(sid) => sid == key_id,
+                    },
+                }
+            };
+            if is_sid_matched || (&envelop.kid == key_id) {
+                information!("Matched envelop: {}", util_envelop::format_envelop(envelop, config));
+                return Ok(envelop);
+            }
+        }
+    }
+
     envelops.iter().enumerate().for_each(|(i, envelop)| {
         println_ex!("#{} {}", i + 1, util_envelop::format_envelop(envelop, config));
     });

src/cmd_info.rs

@@ -46,7 +46,7 @@ pub fn info_single(path: &PathBuf, cmd_info: &CmdInfo) -> XResult<()> {
     let mut file_in = opt_result!(File::open(path), "Open file: {} failed: {}", &path_display);
     let file_in_len = file_in.metadata().map(|m| m.len()).unwrap_or(0);
 
-    let (meta_len, meta) = opt_result!(
+    let (meta_len, _, meta) = opt_result!(
         util_enc_file::read_tiny_encrypt_meta_and_normalize(&mut file_in), "Read file: {}, failed: {}", &path_display
     );
 

src/util_enc_file.rs

@@ -22,13 +22,13 @@ pub fn write_tiny_encrypt_meta(write: &mut impl Write, meta: &TinyEncryptMeta, c
     Ok(encrypted_meta_bytes.len() + 2 + 4)
 }
 
-pub fn read_tiny_encrypt_meta_and_normalize(r: &mut impl Read) -> XResult<(u32, TinyEncryptMeta)> {
+pub fn read_tiny_encrypt_meta_and_normalize(r: &mut impl Read) -> XResult<(u32, bool, TinyEncryptMeta)> {
     let mut meta_len_and_meta = read_tiny_encrypt_meta(r);
-    let _ = meta_len_and_meta.as_mut().map(|ml| ml.1.normalize());
+    let _ = meta_len_and_meta.as_mut().map(|ml| ml.2.normalize());
     meta_len_and_meta
 }
 
-pub fn read_tiny_encrypt_meta(r: &mut impl Read) -> XResult<(u32, TinyEncryptMeta)> {
+pub fn read_tiny_encrypt_meta(r: &mut impl Read) -> XResult<(u32, bool, TinyEncryptMeta)> {
     let mut meta_tag_buff = [0_u8; 2];
     opt_result!(r.read_exact(&mut meta_tag_buff), "Read tag failed: {}");
     let meta_tag = u16::from_be_bytes(meta_tag_buff);
@@ -57,5 +57,5 @@ pub fn read_tiny_encrypt_meta(r: &mut impl Read) -> XResult<(u32, TinyEncryptMet
     }
     debugging!("Encrypted meta: {}", String::from_utf8_lossy(&meta_buff));
 
-    Ok((meta_length, opt_result!(serde_json::from_slice(&meta_buff), "Parse meta failed: {}")))
+    Ok((meta_length, is_meta_compressed, opt_result!(serde_json::from_slice(&meta_buff), "Parse meta failed: {}")))
 }

src/util_env.rs

@@ -8,6 +8,8 @@ use crate::consts;
 pub const TINY_ENCRYPT_ENV_DEFAULT_ALGORITHM: &str = "TINY_ENCRYPT_DEFAULT_ALGORITHM";
 pub const TINY_ENCRYPT_ENV_DEFAULT_COMPRESS: &str = "TINY_ENCRYPT_DEFAULT_COMPRESS";
 pub const TINY_ENCRYPT_ENV_NO_PROGRESS: &str = "TINY_ENCRYPT_NO_PROGRESS";
+pub const TINY_ENCRYPT_ENV_PIN: &str = "TINY_ENCRYPT_PIN";
+pub const TINY_ENCRYPT_ENV_KEY_ID: &str = "TINY_ENCRYPT_KEY_ID";
 
 pub fn get_default_encryption_algorithm() -> Option<&'static str> {
     let env_default_algorithm = env::var(TINY_ENCRYPT_ENV_DEFAULT_ALGORITHM).ok();
@@ -23,6 +25,14 @@ pub fn get_default_encryption_algorithm() -> Option<&'static str> {
     None
 }
 
+pub fn get_pin() -> Option<String> {
+    env::var(TINY_ENCRYPT_ENV_PIN).ok()
+}
+
+pub fn get_key_id() -> Option<String> {
+    env::var(TINY_ENCRYPT_ENV_KEY_ID).ok()
+}
+
 pub fn get_default_compress() -> Option<bool> {
     iff!(rust_util_env::is_env_off(TINY_ENCRYPT_ENV_DEFAULT_COMPRESS), Some(true), None)
 }