feat: v0.1.0, encrypt works

src/cmd_encrypt.rs

@@ -1,15 +1,20 @@
 use std::fs;
+use std::fs::File;
+use std::io::{Read, Write};
 use std::path::PathBuf;
+use std::time::Instant;
 
 use clap::Args;
 use rsa::Pkcs1v15Encrypt;
 use rust_util::{debugging, failure, opt_result, simple_error, success, XResult};
 
+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, TinyEncryptEnvelop, TinyEncryptEnvelopType, TinyEncryptMeta};
-use crate::util::{ENC_AES256_GCM_P256, encode_base64, make_key256_and_nonce, simple_kdf, TINY_ENC_CONFIG_FILE, zeroize};
+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};
 
@@ -18,21 +23,28 @@ pub struct CmdEncrypt {
     /// Files need to be decrypted
     pub paths: Vec<PathBuf>,
     // Comment
+    #[arg(long, short = 'c')]
     pub comment: Option<String>,
     // Comment
+    #[arg(long, short = 'C')]
     pub encrypted_comment: Option<String>,
     // Encryption profile
+    #[arg(long, short = 'p')]
     pub profile: Option<String>,
+    #[arg(long, short = 'x')]
+    pub compress: bool,
 }
 
 pub fn encrypt(cmd_encrypt: CmdEncrypt) -> XResult<()> {
     let config = TinyEncryptConfig::load(TINY_ENC_CONFIG_FILE)?;
+    debugging!("Found tiny encrypt config: {:?}", config);
     let envelops = config.find_envelops(&cmd_encrypt.profile)?;
     if envelops.is_empty() { return simple_error!("Cannot find any valid envelops"); }
+    debugging!("Found envelops: {:?}", envelops);
 
     debugging!("Cmd encrypt: {:?}", cmd_encrypt);
     for path in &cmd_encrypt.paths {
-        match encrypt_single(path, &envelops) {
+        match encrypt_single(path, &envelops, &cmd_encrypt) {
             Ok(_) => success!("Encrypt {} succeed", path.to_str().unwrap_or("N/A")),
             Err(e) => failure!("Encrypt {} failed: {}", path.to_str().unwrap_or("N/A"), e),
         }
@@ -40,27 +52,92 @@ pub fn encrypt(cmd_encrypt: CmdEncrypt) -> XResult<()> {
     Ok(())
 }
 
-fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop]) -> XResult<()> {
+fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_encrypt: &CmdEncrypt) -> XResult<()> {
+    let path_display = format!("{}", path.display());
+    util::require_none_tiny_enc_file_and_exists(path)?;
+
+    let mut file_in = opt_result!(File::open(path), "Open file: {} failed: {}", &path_display);
+
+    let path_out = format!("{}{}", path_display, util::TINY_ENC_FILE_EXT);
+    util::require_file_not_exists(path_out.as_str())?;
+
     let (key, nonce) = make_key256_and_nonce();
     let envelops = encrypt_envelops(&key, &envelops)?;
 
+    let encrypted_comment = match &cmd_encrypt.encrypted_comment {
+        None => None,
+        Some(encrypted_comment) => Some(encode_base64(
+            &aes_gcm_encrypt(&key, &nonce, encrypted_comment.as_bytes())?))
+    };
+
     let file_metadata = opt_result!(fs::metadata(path), "Read file: {} meta failed: {}", path.display());
     let enc_metadata = EncMetadata {
-        comment: None,
-        encrypted_comment: None,
+        comment: cmd_encrypt.comment.clone(),
+        encrypted_comment,
         encrypted_meta: None,
-        compress: false,
+        compress: cmd_encrypt.compress,
     };
 
     let encrypt_meta = TinyEncryptMeta::new(&file_metadata, &enc_metadata, &nonce, envelops);
     debugging!("Encrypted meta: {:?}", encrypt_meta);
 
-    // TODO write to file and do encrypt
+
+    let mut file_out = File::create(&path_out)?;
+    opt_result!(file_out.write_all(&util::TINY_ENC_MAGIC_TAG.to_be_bytes()), "Write tag failed: {}");
+    let encrypted_meta_bytes = opt_result!(serde_json::to_vec(&encrypt_meta), "Generate meta json bytes failed: {}");
+    let encrypted_meta_bytes_len = encrypted_meta_bytes.len() as u32;
+    opt_result!(file_out.write_all(&encrypted_meta_bytes_len.to_be_bytes()), "Write meta len failed: {}");
+    opt_result!(file_out.write_all(&encrypted_meta_bytes), "Write meta failed: {}");
+
+    let start = Instant::now();
+    encrypt_file(&mut file_in, &mut file_out, &key, &nonce, cmd_encrypt.compress)?;
+    let encrypt_duration = start.elapsed();
+    debugging!("Encrypt file: {} elapsed: {} ms", path_display, encrypt_duration.as_millis());
+
     zeroize(key);
     zeroize(nonce);
     Ok(())
 }
 
+
+fn encrypt_file(file_in: &mut File, file_out: &mut File, key: &[u8], nonce: &[u8], compress: bool) -> XResult<usize> {
+    let mut total_len = 0;
+    let mut buffer = [0u8; 1024 * 8];
+    let key = opt_result!(key.try_into(), "Key is not 32 bytes: {}");
+    let mut gz_encoder = GzStreamEncoder::new_default();
+    let mut encryptor = aes_gcm_stream::Aes256GcmStreamEncryptor::new(key, &nonce);
+    loop {
+        let len = opt_result!(file_in.read(&mut buffer), "Read file failed: {}");
+        if len == 0 {
+            let last_block = if compress {
+                let last_compressed_buffer = opt_result!(gz_encoder.finalize(), "Decompress file failed: {}");
+                let mut encrypted_block = encryptor.update(&last_compressed_buffer);
+                let (last_block, tag) = encryptor.finalize();
+                encrypted_block.extend_from_slice(&last_block);
+                encrypted_block.extend_from_slice(&tag);
+                encrypted_block
+            } else {
+                let (mut last_block, tag) = encryptor.finalize();
+                last_block.extend_from_slice(&tag);
+                last_block
+            };
+            opt_result!(file_out.write_all(&last_block), "Write file failed: {}");
+            success!("Decrypt finished, total bytes: {}", total_len);
+            break;
+        } else {
+            total_len += len;
+            let encrypted = if compress {
+                let compressed = opt_result!(gz_encoder.update(&buffer[0..len]), "Decompress file failed: {}");
+                encryptor.update(&compressed)
+            } else {
+                encryptor.update(&buffer[0..len])
+            };
+            opt_result!(file_out.write_all(&encrypted), "Write file failed: {}");
+        }
+    }
+    Ok(total_len)
+}
+
 fn encrypt_envelops(key: &[u8], envelops: &[&TinyEncryptConfigEnvelop]) -> XResult<Vec<TinyEncryptEnvelop>> {
     let mut encrypted_envelops = vec![];
     for envelop in envelops {
@@ -89,7 +166,7 @@ fn encrypt_envelop_ecdh(key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResu
         header: WrapKeyHeader {
             kid: Some(envelop.kid.clone()),
             enc: ENC_AES256_GCM_P256.to_string(),
-            e_pub_key: encode_base64(&ephemeral_spki),
+            e_pub_key: encode_base64_url_no_pad(&ephemeral_spki),
         },
         nonce,
         encrypted_data: encrypted_key,