hatter/tiny-encrypt-rs
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

feat: improve decrypt

Browse Source
This commit is contained in:
Hatter Jiang
2023-12-01 21:52:26 +08:00
parent 4966ddf72b
commit ae84bbd13b
6 changed files with 105 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.lock generated
View File

@@ -1657,7 +1657,7 @@ dependencies = [
[[package]]
name = "tiny-encrypt"
version = "0.8.0"
version = "0.8.1"
dependencies = [
"aes-gcm-stream",
"base64",

6
Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "tiny-encrypt"
version = "0.8.0"
version = "0.8.1"
edition = "2021"
license = "MIT"
description = "A simple and tiny file encrypt tool"
@@ -28,7 +28,7 @@ p256 = { version = "0.13", features = ["pem", "ecdh", "pkcs8"] }
p384 = { version = "0.13", features = ["pem", "ecdh", "pkcs8"] }
rand = "0.8"
# reqwest = { version = "0.11", features = ["blocking", "rustls", "rustls-tls"] }
rpassword = "7.2"
rpassword = "7.3"
rsa = { version = "0.9", features = ["pem"] }
rust-crypto = "0.2"
rust_util = "0.6"
@@ -39,7 +39,7 @@ tabled = "0.14"
x25519-dalek = "2.0"
x509-parser = "0.15"
yubikey = { version = "0.8", features = ["untested"], optional = true }
zeroize = "1.6"
zeroize = "1.7"
[patch.crates-io]
rust-crypto = { git = "https://github.com/jht5945/rust-crypto.git" }

123
src/cmd_decrypt.rs
View File

@@ -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));
});

2
src/cmd_info.rs
View File

@@ -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
);

8
src/util_enc_file.rs
View File

@@ -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: {}")))
}

10
src/util_env.rs
View File

@@ -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)
}