From 71ebe68a860138cb2cb33c309e22a276b4eaf712 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 30 Sep 2023 21:40:27 +0800
Subject: [PATCH] feat: supports -R

---
 src/cmd_decrypt.rs | 27 +++++++++++++++++++--------
 src/cmd_encrypt.rs | 15 +++++++++++++--
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/cmd_decrypt.rs b/src/cmd_decrypt.rs
index 12b8787..d222129 100644
--- a/src/cmd_decrypt.rs
+++ b/src/cmd_decrypt.rs
@@ -1,5 +1,5 @@
+use std::{fs, io};
 use std::fs::File;
-use std::io;
 use std::io::{Read, Write};
 use std::path::PathBuf;
 use std::str::FromStr;
@@ -8,7 +8,7 @@ use std::time::Instant;
 use clap::Args;
 use openpgp_card::crypto_data::Cryptogram;
 use openpgp_card::OpenPgp;
-use rust_util::{debugging, failure, opt_result, simple_error, success, util_term, XResult};
+use rust_util::{debugging, failure, information, opt_result, simple_error, success, util_term, warning, XResult};
 use x509_parser::prelude::FromDer;
 use x509_parser::x509::SubjectPublicKeyInfo;
 use yubikey::piv::{AlgorithmId, decrypt_data, RetiredSlotId, SlotId};
@@ -32,12 +32,15 @@ pub struct CmdDecrypt {
     /// Slot
     #[arg(long, short = 's')]
     pub slot: Option<String>,
+    /// Remove source file
+    #[arg(long, short = 'R')]
+    pub remove_source_file: bool,
 }
 
 pub fn decrypt(cmd_decrypt: CmdDecrypt) -> XResult<()> {
     debugging!("Cmd decrypt: {:?}", cmd_decrypt);
     for path in &cmd_decrypt.paths {
-        match decrypt_single(path, &cmd_decrypt.pin, &cmd_decrypt.slot) {
+        match decrypt_single(path, &cmd_decrypt.pin, &cmd_decrypt.slot, &cmd_decrypt) {
             Ok(_) => success!("Decrypt {} succeed", path.to_str().unwrap_or("N/A")),
             Err(e) => failure!("Decrypt {} failed: {}", path.to_str().unwrap_or("N/A"), e),
         }
@@ -45,7 +48,7 @@ pub fn decrypt(cmd_decrypt: CmdDecrypt) -> XResult<()> {
     Ok(())
 }
 
-pub fn decrypt_single(path: &PathBuf, pin: &Option<String>, slot: &Option<String>) -> XResult<()> {
+pub fn decrypt_single(path: &PathBuf, pin: &Option<String>, slot: &Option<String>, cmd_decrypt: &CmdDecrypt) -> XResult<()> {
     let path_display = format!("{}", path.display());
     util::require_tiny_enc_file_and_exists(path)?;
 
@@ -73,6 +76,14 @@ pub fn decrypt_single(path: &PathBuf, pin: &Option<String>, slot: &Option<String
 
     util::zeroize(key);
     util::zeroize(nonce);
+    drop(file_in);
+    drop(file_out);
+    if cmd_decrypt.remove_source_file {
+        match fs::remove_file(path) {
+            Err(e) => warning!("Remove file: {} failed: {}", path_display, e),
+            Ok(_) => information!("Remove file: {} succeed", path_display),
+        }
+    }
     Ok(())
 }
 
@@ -141,7 +152,7 @@ fn try_decrypt_key_ecdh(envelop: &TinyEncryptEnvelop, pin: &Option<String>, slot
                 &epk_bytes,
                 AlgorithmId::EccP256,
                 slot_id,
-            ), "Decrypt piv failed: {}");
+            ), "Decrypt via PIV card failed: {}");
     let key = simple_kdf(decrypted_shared_secret.as_slice());
     let decrypted_key = aes_gcm_decrypt(&key, &wrap_key.nonce, &wrap_key.encrypted_data)?;
     Ok(decrypted_key)
@@ -181,10 +192,10 @@ fn read_slot(slot: &Option<String>) -> XResult<String> {
             io::stdout().flush().ok();
             let mut buff = String::new();
             let _ = io::stdin().read_line(&mut buff).expect("Read line from stdin");
-            if buff.is_empty() {
-                simple_error!("Slot is required, and not inputed")
+            if buff.trim().is_empty() {
+                simple_error!("Slot is required, and not inputted")
             } else {
-                Ok(buff)
+                Ok(buff.trim().to_string())
             }
         }
     }
diff --git a/src/cmd_encrypt.rs b/src/cmd_encrypt.rs
index b7a31e6..62c62ec 100644
--- a/src/cmd_encrypt.rs
+++ b/src/cmd_encrypt.rs
@@ -6,14 +6,14 @@ use std::time::Instant;
 
 use clap::Args;
 use rsa::Pkcs1v15Encrypt;
-use rust_util::{debugging, failure, opt_result, simple_error, success, XResult};
+use rust_util::{debugging, failure, information, opt_result, simple_error, success, warning, XResult};
 
-use crate::util;
 use crate::compress::GzStreamEncoder;
 use crate::config::{TinyEncryptConfig, TinyEncryptConfigEnvelop};
 use crate::crypto_aes::aes_gcm_encrypt;
 use crate::crypto_rsa::parse_spki;
 use crate::spec::{EncMetadata, TINY_ENCRYPT_VERSION_10, TinyEncryptEnvelop, TinyEncryptEnvelopType, TinyEncryptMeta};
+use crate::util;
 use crate::util::{ENC_AES256_GCM_P256, encode_base64, encode_base64_url_no_pad, make_key256_and_nonce, simple_kdf, TINY_ENC_CONFIG_FILE, zeroize};
 use crate::util_ecdh::compute_shared_secret;
 use crate::wrap_key::{WrapKey, WrapKeyHeader};
@@ -37,6 +37,9 @@ pub struct CmdEncrypt {
     /// Compatible with 1.0
     #[arg(long, short = '1')]
     pub compatible_with_1_0: bool,
+    /// Remove source file
+    #[arg(long, short = 'R')]
+    pub remove_source_file: bool,
 }
 
 pub fn encrypt(cmd_encrypt: CmdEncrypt) -> XResult<()> {
@@ -120,6 +123,14 @@ fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_en
 
     zeroize(key);
     zeroize(nonce);
+    drop(file_in);
+    drop(file_out);
+    if cmd_encrypt.remove_source_file {
+        match fs::remove_file(path) {
+            Err(e) => warning!("Remove file: {} failed: {}", path_display, e),
+            Ok(_) => information!("Remove file: {} succeed", path_display),
+        }
+    }
     Ok(())
 }