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

feat: optimize code

Browse Source
This commit is contained in:
Hatter Jiang
2023-10-15 11:53:29 +08:00
parent 12aa3b48d2
commit 1a3c670bad
5 changed files with 52 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
src/card.rs
View File

@@ -1,15 +0,0 @@
use openpgp_card_pcsc::PcscBackend;
use rust_util::{opt_result, opt_value_result, simple_error, warning, XResult};
pub fn get_card() -> XResult<PcscBackend> {
let card_list = opt_result!(
PcscBackend::cards(None), "Read OpenPGP card list failed: {}"
);
if card_list.is_empty() {
return simple_error!("Cannot find any card");
}
if card_list.len() > 1 {
warning!("Find {} OpenPGP cards, will use first card", card_list.len());
}
Ok(opt_value_result!(card_list.into_iter().next(), "SHOULD NOT HAPPEN, CANNOT FIND ANY CARD"))
}

73
src/cmd_decrypt.rs
View File

@@ -1,16 +1,14 @@
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::time::{Instant, SystemTime};
use clap::Args;
use fs_set_times::SystemTimeSpec;
use openpgp_card::{OpenPgp, OpenPgpTransaction};
use openpgp_card::crypto_data::Cryptogram;
use rust_util::{
debugging, failure, iff, information, opt_result, simple_error, success,
util_msg, util_term, warning, XResult,
util_msg, warning, XResult,
};
use rust_util::util_time::UnixEpochTime;
use x509_parser::prelude::FromDer;
@@ -19,7 +17,7 @@ use yubikey::piv::{AlgorithmId, decrypt_data};
use yubikey::YubiKey;
use zeroize::Zeroize;
use crate::{card, file, util, util_piv};
use crate::{file, util, util_pgp, util_piv};
use crate::compress::GzStreamDecoder;
use crate::config::TinyEncryptConfig;
use crate::consts::{DATE_TIME_FORMAT, ENC_AES256_GCM_P256, ENC_AES256_GCM_P384, ENC_AES256_GCM_X25519, SALT_COMMENT, TINY_ENC_CONFIG_FILE, TINY_ENC_FILE_EXT};
@@ -256,8 +254,8 @@ fn try_decrypt_key_ecdh(config: &Option<TinyEncryptConfig>,
let e_pub_key_bytes = opt_result!(util::decode_base64_url_no_pad(e_pub_key), "Invalid envelop: {}");
let (_, subject_public_key_info) = opt_result!(SubjectPublicKeyInfo::from_der(&e_pub_key_bytes), "Invalid envelop: {}");
let slot = read_slot(config, &envelop.kid, slot)?;
let pin = read_pin(pin);
let slot = util_piv::read_piv_slot(config, &envelop.kid, slot)?;
let pin = util::read_pin(pin);
let epk_bytes = subject_public_key_info.subject_public_key.as_ref();
let mut yk = opt_result!(YubiKey::open(), "YubiKey not found: {}");
@@ -287,10 +285,10 @@ fn try_decrypt_key_ecdh_pgp_x25519(envelop: &TinyEncryptEnvelop, pin: &Option<St
let e_pub_key = &wrap_key.header.e_pub_key;
let epk_bytes = opt_result!(util::decode_base64_url_no_pad(e_pub_key), "Invalid envelop: {}");
let mut pgp = get_openpgp()?;
let mut pgp = util_pgp::get_openpgp()?;
let mut trans = opt_result!(pgp.transaction(), "Open card failed: {}");
read_and_verify_openpgp_pin(&mut trans, pin)?;
util_pgp::read_and_verify_openpgp_pin(&mut trans, pin)?;
let shared_secret = trans.decipher(Cryptogram::ECDH(&epk_bytes))?;
@@ -302,10 +300,10 @@ fn try_decrypt_key_ecdh_pgp_x25519(envelop: &TinyEncryptEnvelop, pin: &Option<St
}
fn try_decrypt_key_pgp(envelop: &TinyEncryptEnvelop, pin: &Option<String>) -> XResult<Vec<u8>> {
let mut pgp = get_openpgp()?;
let mut pgp = util_pgp::get_openpgp()?;
let mut trans = opt_result!(pgp.transaction(), "Open card failed: {}");
read_and_verify_openpgp_pin(&mut trans, pin)?;
util_pgp::read_and_verify_openpgp_pin(&mut trans, pin)?;
let pgp_envelop = &envelop.encrypted_key;
debugging!("PGP envelop: {}", &pgp_envelop);
@@ -315,61 +313,6 @@ fn try_decrypt_key_pgp(envelop: &TinyEncryptEnvelop, pin: &Option<String>) -> XR
Ok(key)
}
fn read_and_verify_openpgp_pin(trans: &mut OpenPgpTransaction, pin: &Option<String>) -> XResult<()> {
let pin = read_pin(pin);
if let Err(e) = trans.verify_pw1_user(pin.as_ref()) {
failure!("Verify user pin failed: {}", e);
return simple_error!("User pin verify failed: {}", e);
}
success!("User pin verify success!");
Ok(())
}
fn get_openpgp() -> XResult<OpenPgp> {
let card = match card::get_card() {
Err(e) => {
failure!("Get PGP card failed: {}", e);
return simple_error!("Get card failed: {}", e);
}
Ok(card) => card
};
Ok(OpenPgp::new(card))
}
fn read_slot(config: &Option<TinyEncryptConfig>, kid: &str, slot: &Option<String>) -> XResult<String> {
match slot {
Some(slot) => Ok(slot.to_string()),
None => {
if let Some(config) = config {
if let Some(first_arg) = config.find_first_arg_by_kid(kid) {
information!("Found kid: {}'s slot: {}", kid, first_arg);
return Ok(first_arg.to_string());
}
}
print!("Input slot(eg 82, 83 ...): ");
io::stdout().flush().ok();
let mut buff = String::new();
let _ = io::stdin().read_line(&mut buff).expect("Read line from stdin");
if buff.trim().is_empty() {
simple_error!("Slot is required, and not inputted")
} else {
Ok(buff.trim().to_string())
}
}
}
}
fn read_pin(pin: &Option<String>) -> String {
match pin {
Some(pin) => pin.to_string(),
None => if util_term::read_yes_no("Use default PIN 123456, please confirm") {
"123456".into()
} else {
rpassword::prompt_password("Please input PIN: ").expect("Read PIN failed")
}
}
}
fn select_envelop(meta: &TinyEncryptMeta) -> XResult<&TinyEncryptEnvelop> {
let envelops = match &meta.envelops {
None => return simple_error!("No envelops found"),

4
src/main.rs
View File

@@ -2,8 +2,8 @@ extern crate core;
use clap::{Parser, Subcommand};
use rust_util::XResult;
use crate::cmd_config::CmdConfig;
use crate::cmd_config::CmdConfig;
use crate::cmd_decrypt::CmdDecrypt;
use crate::cmd_encrypt::CmdEncrypt;
use crate::cmd_info::CmdInfo;
@@ -12,6 +12,7 @@ use crate::cmd_version::CmdVersion;
mod consts;
mod util;
mod util_piv;
mod util_pgp;
mod util_ecdh;
mod util_p384;
mod util_x25519;
@@ -22,7 +23,6 @@ mod crypto_aes;
mod crypto_rsa;
mod wrap_key;
mod file;
mod card;
mod cmd_version;
mod cmd_config;
mod cmd_info;

13
src/util.rs
View File

@@ -5,11 +5,22 @@ use std::path::{Path, PathBuf};
use base64::Engine;
use base64::engine::general_purpose;
use rand::random;
use rust_util::{information, simple_error, warning, XResult};
use rust_util::{information, simple_error, util_term, warning, XResult};
use zeroize::Zeroize;
use crate::consts::TINY_ENC_FILE_EXT;
pub fn read_pin(pin: &Option<String>) -> String {
match pin {
Some(pin) => pin.to_string(),
None => if util_term::read_yes_no("Use default PIN 123456, please confirm") {
"123456".into()
} else {
rpassword::prompt_password("Please input PIN: ").expect("Read PIN failed")
}
}
}
pub fn remove_file_with_msg(path: &PathBuf) {
match fs::remove_file(path) {
Err(e) => warning!("Remove file: {} failed: {}", path.display(), e),

31
src/util_piv.rs
View File

@@ -1,6 +1,35 @@
use rust_util::{simple_error, XResult};
use std::io;
use std::io::Write;
use rust_util::{information, simple_error, XResult};
use yubikey::piv::{RetiredSlotId, SlotId};
use crate::config::TinyEncryptConfig;
pub fn read_piv_slot(config: &Option<TinyEncryptConfig>, kid: &str, slot: &Option<String>) -> XResult<String> {
match slot {
Some(slot) => Ok(slot.to_string()),
None => {
if let Some(config) = config {
if let Some(first_arg) = config.find_first_arg_by_kid(kid) {
information!("Found kid: {}'s slot: {}", kid, first_arg);
return Ok(first_arg.to_string());
}
}
print!("Input slot(eg 82, 83 ...): ");
io::stdout().flush().ok();
let mut buff = String::new();
let _ = io::stdin().read_line(&mut buff).expect("Read line from stdin");
if buff.trim().is_empty() {
simple_error!("Slot is required, and not inputted")
} else {
Ok(buff.trim().to_string())
}
}
}
}
pub fn get_slot_id(slot: &str) -> XResult<SlotId> {
let slot_lower = slot.to_lowercase();
Ok(match slot_lower.as_str() {