tiny-encrypt 1.9.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 flate2::Compression;
use rsa::Pkcs1v15Encrypt;
use rust_util::util_time::UnixEpochTime;
use rust_util::{debugging, failure, iff, information, opt_result, simple_error, success, util_size, warning, XResult};

use crate::compress::GzStreamEncoder;
use crate::config::{TinyEncryptConfig, TinyEncryptConfigEnvelop};
use crate::consts::{ENC_AES256_GCM_KYBER1204, ENC_AES256_GCM_MLKEM1024, ENC_AES256_GCM_MLKEM768, ENC_AES256_GCM_P256, ENC_AES256_GCM_P384, ENC_AES256_GCM_X25519, ENC_CHACHA20_POLY1305_KYBER1204, ENC_CHACHA20_POLY1305_MLKEM1024, ENC_CHACHA20_POLY1305_MLKEM768, ENC_CHACHA20_POLY1305_P256, ENC_CHACHA20_POLY1305_P384, ENC_CHACHA20_POLY1305_X25519, SALT_COMMENT, TINY_ENC_FILE_EXT, TINY_ENC_PEM_FILE_EXT, TINY_ENC_PEM_NAME};
use crate::crypto_cryptor::{Cryptor, KeyNonce};
use crate::spec::{
    EncEncryptedMeta, EncMetadata,
    TinyEncryptEnvelop, TinyEncryptEnvelopType, TinyEncryptMeta,
};
use crate::util::{decode_base64, is_tiny_enc_file, to_pem};
use crate::util_ecdh::{ecdh_kyber1024, ecdh_p256, ecdh_p384, ecdh_x25519};
use crate::util_progress::Progress;
use crate::{util_mlkem, util_rsa};
use crate::wrap_key::{WrapKey, WrapKeyHeader};
use crate::{crypto_cryptor, crypto_simple, util, util_enc_file, util_env, util_gpg};
use crate::temporary_key::parse_temporary_keys;
use crate::util_mlkem::MlKemAlgo;

#[derive(Debug, Args)]
pub struct CmdEncrypt {
    /// Plaintext comment
    #[arg(long, short = 'c')]
    pub comment: Option<String>,

    /// Encrypted comment
    #[arg(long, short = 'C')]
    pub encrypted_comment: Option<String>,

    /// Encryption profile (use default when --key-filter is assigned)
    #[arg(long, short = 'p')]
    pub profile: Option<String>,

    /// Encryption key filter (key_id or type:TYPE(e.g. ecdh, pgp, ecdh-p384, pgp-ed25519), multiple joined by ',', ALL for all)
    #[arg(long, short = 'k')]
    pub key_filter: Option<String>,

    /// Temporary key
    #[arg(long)]
    pub temporary_key: Option<Vec<String>>,

    /// Compress before encrypt
    #[arg(long, short = 'x')]
    pub compress: bool,

    /// Compress level (from 0[none], 1[fast] .. 6[default] .. to 9[best])
    #[arg(long, short = 'L')]
    pub compress_level: Option<u32>,

    /// Remove source file
    #[arg(long, short = 'R')]
    pub remove_file: bool,

    /// Create file (create a empty encrypted file)
    #[arg(long, short = 'a')]
    pub create: bool,

    /// Disable compress meta
    #[arg(long)]
    pub disable_compress_meta: bool,

    /// Output file in PEM format (alias --pem)
    #[arg(long, alias = "pem")]
    pub pem_output: bool,

    /// Encryption algorithm (AES/GCM, CHACHA20/POLY1305 or AES, CHACHA20, default AES/GCM)
    #[arg(long, short = 'A')]
    pub encryption_algorithm: Option<String>,

    /// Config file or based64 encoded (starts with: base64:)
    #[arg(long)]
    pub config: Option<String>,

    /// Files need to be decrypted
    pub paths: Vec<PathBuf>,
}

pub fn encrypt(cmd_encrypt: CmdEncrypt) -> XResult<()> {
    let config = TinyEncryptConfig::load_default(&cmd_encrypt.config)?;
    debugging!("Found tiny encrypt config: {:?}", config);
    let mut envelops = config.find_envelops(&cmd_encrypt.profile, &cmd_encrypt.key_filter)?;
    debugging!("Found envelops: {:?}", envelops);

    let temporary_envelops = parse_temporary_keys(&cmd_encrypt.temporary_key)?;
    if !temporary_envelops.is_empty() {
        for t_envelop in &temporary_envelops {
            envelops.push(t_envelop)
        }
        debugging!("Final envelops: {:?}", envelops);
    }
    if envelops.is_empty() { return simple_error!("Cannot find any valid envelops"); }

    let envelop_tkids: Vec<_> = envelops.iter()
        .map(|e| format!("{}:{}", e.r#type.get_name(), e.kid))
        .collect();
    information!("Matched {} envelop(s): \n- {}", envelops.len(), envelop_tkids.join("\n- "));

    debugging!("Cmd encrypt: {:?}", cmd_encrypt);
    let start = Instant::now();
    let mut succeed_count = 0;
    let mut skipped_count = 0;
    let mut failed_count = 0;
    let mut total_len = 0_u64;
    for path in &cmd_encrypt.paths {
        let path = config.resolve_path_namespace(path, false);
        let start_encrypt_single = Instant::now();
        match encrypt_single(&path, &envelops, &cmd_encrypt) {
            Ok(len) => {
                total_len += len;
                if len > 0 { succeed_count += 1; } else { skipped_count += 1; }
                success!(
                    "Encrypt {} succeed, cost {} ms, file size {} byte(s)",
                    path.to_str().unwrap_or("N/A"),
                    start_encrypt_single.elapsed().as_millis(),
                    len
                );
            }
            Err(e) => {
                failed_count += 1;
                failure!("Encrypt {} failed: {}", path.to_str().unwrap_or("N/A"), e);
            }
        }
    }
    if (succeed_count + failed_count) > 1 {
        success!(
            "Encrypt succeed {} file(s) {} byte(s), failed {} file(s), skipped {} file(s), total cost {} ms",
            succeed_count,
            total_len,
            failed_count,
            skipped_count,
            start.elapsed().as_millis(),
        );
    }
    Ok(())
}

pub fn encrypt_single(path: &PathBuf, envelops: &[&TinyEncryptConfigEnvelop], cmd_encrypt: &CmdEncrypt) -> XResult<u64> {
    let path_display = format!("{}", path.display());
    let path_out = if cmd_encrypt.pem_output {
        format!("{}{}", path_display, TINY_ENC_PEM_FILE_EXT)
    } else {
        format!("{}{}", path_display, TINY_ENC_FILE_EXT)
    };
    let encrypt_single_result = encrypt_single_file_out(path, &path_out, envelops, cmd_encrypt);
    if cmd_encrypt.create {
        if let Ok(content) = fs::read_to_string(path) {
            if content == "\n" {
                let _ = fs::remove_file(path);
            }
        }
    }
    encrypt_single_result
}

pub fn encrypt_single_file_out(path: &PathBuf, path_out: &str, envelops: &[&TinyEncryptConfigEnvelop], cmd_encrypt: &CmdEncrypt) -> XResult<u64> {
    let path_display = format!("{}", path.display());
    if is_tiny_enc_file(&path_display) {
        information!("Tiny enc file skipped: {}", path_display);
        return Ok(0);
    }

    let cryptor = crypto_cryptor::get_cryptor_by_encryption_algorithm(&cmd_encrypt.encryption_algorithm)?;
    information!("Using encryption algorithm: {}", cryptor.get_name());

    if cmd_encrypt.create {
        util::require_file_not_exists(path)?;
        opt_result!(fs::write(path, "\n"), "Write empty file failed: {}");
    } else {
        util::require_file_exists(path)?;
    }

    let mut file_in = opt_result!(File::open(path), "Open file: {} failed: {}", &path_display);

    util::require_file_not_exists(path_out)?;

    let (key, nonce) = util::make_key256_and_nonce();
    let key_nonce = KeyNonce { k: key.as_ref(), n: nonce.as_ref() };
    // Encrypt session key to envelops
    let envelops = encrypt_envelops(cryptor, key.as_ref(), envelops)?;

    let encrypted_comment = match &cmd_encrypt.encrypted_comment {
        None => None,
        Some(encrypted_comment) => Some(util::encode_base64(
            &crypto_simple::encrypt_with_salt(
                cryptor, &key_nonce, SALT_COMMENT, encrypted_comment.as_bytes())?))
    };

    let file_metadata = opt_result!(fs::metadata(path), "Read file: {} meta failed: {}", path.display());
    let enc_encrypted_meta = EncEncryptedMeta {
        filename: Some(util::get_file_name(path)),
        c_time: file_metadata.created().ok().and_then(|t| t.to_millis()),
        m_time: file_metadata.modified().ok().and_then(|t| t.to_millis()),
    };
    let enc_encrypted_meta_bytes = opt_result!(enc_encrypted_meta.seal(
        cryptor, &key_nonce), "Seal enc-encrypted-meta failed: {}");

    let compress_level = get_compress_level(cmd_encrypt, &path_display, cmd_encrypt.pem_output);
    debugging!("Compress level: {:?}", compress_level);

    let enc_metadata = EncMetadata {
        comment: cmd_encrypt.comment.clone(),
        encrypted_comment,
        encrypted_meta: Some(util::encode_base64(&enc_encrypted_meta_bytes)),
        compress: compress_level.is_some(),
    };

    let encrypt_meta = TinyEncryptMeta::new(
        &file_metadata, &enc_metadata, cryptor, nonce.as_ref(), envelops);
    debugging!("Encrypted meta: {:?}", encrypt_meta);

    let start = Instant::now();

    let mut file_out = File::create(path_out)?;
    let compress_meta = !cmd_encrypt.disable_compress_meta;

    if cmd_encrypt.pem_output {
        let temp_output_len = file_in.metadata().map(|m| m.len()).unwrap_or(0) + 1024 * 8;
        if temp_output_len > 8 * 1024 * 1028 {
            warning!("Input file is more than 8 MiB.");
        }
        if temp_output_len > 32 * 1024 * 1028 {
            return simple_error!("Input file is too large, file is {} bytes", temp_output_len);
        }
        let mut temp_output = Vec::with_capacity(temp_output_len as usize);
        let _ = util_enc_file::write_tiny_encrypt_meta(&mut temp_output, &encrypt_meta, compress_meta)?;
        encrypt_file(&mut file_in, file_metadata.len(), &mut temp_output,
                     cryptor, &key_nonce, &compress_level,
        )?;
        let temp_output_pem = to_pem(&temp_output, TINY_ENC_PEM_NAME);
        file_out.write_all(temp_output_pem.as_bytes())?;
    } else {
        let _ = util_enc_file::write_tiny_encrypt_meta(&mut file_out, &encrypt_meta, compress_meta)?;
        encrypt_file(&mut file_in, file_metadata.len(), &mut file_out,
                     cryptor, &key_nonce, &compress_level,
        )?;
    }

    drop(file_out);
    let encrypt_duration = start.elapsed();
    let compress_desc = iff!(compress_level.is_some(), " [with compress]", "");
    debugging!("Inner encrypt file{}: {} elapsed: {} ms", compress_desc, path_display, encrypt_duration.as_millis());
    if cmd_encrypt.remove_file { util::remove_file_with_msg(path); }
    Ok(file_metadata.len())
}

pub(crate) fn encrypt_file(file_in: &mut impl Read, file_len: u64, file_out: &mut impl Write, cryptor: Cryptor,
                           key_nonce: &KeyNonce, compress_level: &Option<u32>) -> XResult<u64> {
    let compress = compress_level.is_some();
    let mut total_len = 0_u64;
    let mut write_len = 0_u64;
    let mut buffer = [0u8; 1024 * 8];
    let mut gz_encoder = match compress_level {
        None => GzStreamEncoder::new_default(),
        Some(compress_level) => {
            if *compress_level > 9 {
                return simple_error!("Compress level must be in range [0, 9]");
            }
            GzStreamEncoder::new(Compression::new(*compress_level))
        }
    };
    let progress = Progress::new(file_len);
    let mut encryptor = cryptor.encryptor(key_nonce)?;
    loop {
        let len = opt_result!(file_in.read(&mut buffer), "Read file failed: {}");
        if len == 0 {
            let last_block_and_tag = 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();
                write_len += encrypted_block.len() as u64;
                write_len += last_block.len() as u64;
                encrypted_block.extend_from_slice(&last_block);
                encrypted_block.extend_from_slice(&tag);
                encrypted_block
            } else {
                let (mut last_block, tag) = encryptor.finalize();
                write_len += last_block.len() as u64;
                last_block.extend_from_slice(&tag);
                last_block
            };
            opt_result!(file_out.write_all(&last_block_and_tag), "Write file failed: {}");
            progress.finish();
            debugging!("Encrypt finished, total bytes: {} byte(s)", total_len);
            if compress {
                information!("File is compressed: {} -> {}, ratio: {}%",
                    util_size::get_display_size(total_len as i64),
                    util_size::get_display_size(write_len as i64),
                    util::ratio(write_len, total_len));
            }
            break;
        } else {
            total_len += len as u64;
            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])
            };
            write_len += encrypted.len() as u64;
            opt_result!(file_out.write_all(&encrypted), "Write file failed: {}");
            progress.position(total_len);
        }
    }
    Ok(total_len)
}

pub fn encrypt_envelops(cryptor: Cryptor, key: &[u8], envelops: &[&TinyEncryptConfigEnvelop]) -> XResult<Vec<TinyEncryptEnvelop>> {
    let mut encrypted_envelops = vec![];
    for envelop in envelops {
        match envelop.r#type {
            TinyEncryptEnvelopType::PgpRsa | TinyEncryptEnvelopType::PivRsa => {
                encrypted_envelops.push(encrypt_envelop_rsa(key, envelop)?);
            }
            TinyEncryptEnvelopType::Gpg => {
                encrypted_envelops.push(encrypt_envelop_gpg(key, envelop)?);
            }
            TinyEncryptEnvelopType::PgpX25519 | TinyEncryptEnvelopType::StaticX25519 => {
                encrypted_envelops.push(encrypt_envelop_ecdh_x25519(cryptor, key, envelop)?);
            }
            TinyEncryptEnvelopType::PivP256 | TinyEncryptEnvelopType::KeyP256 | TinyEncryptEnvelopType::ExtP256 => {
                encrypted_envelops.push(encrypt_envelop_ecdh_p256(cryptor, key, envelop)?);
            }
            TinyEncryptEnvelopType::PivP384 | TinyEncryptEnvelopType::ExtP384 => {
                encrypted_envelops.push(encrypt_envelop_ecdh_p384(cryptor, key, envelop)?);
            }
            TinyEncryptEnvelopType::StaticKyber1024 => {
                encrypted_envelops.push(encrypt_envelop_ecdh_kyber1204(cryptor, key, envelop)?);
            }
            TinyEncryptEnvelopType::KeyMlKem768 | TinyEncryptEnvelopType::KeyMlKem1024 |
            TinyEncryptEnvelopType::ExtMlKem768 | TinyEncryptEnvelopType::ExtMlKem1024 => {
                encrypted_envelops.push(encrypt_envelop_ecdh_ml_kem(cryptor, key, envelop)?);
            }
            _ => return simple_error!("Not supported type: {:?}", envelop.r#type),
        }
    }
    Ok(encrypted_envelops)
}

fn encrypt_envelop_ecdh_p256(cryptor: Cryptor, key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let public_key_point_hex = &envelop.public_part;
    let (shared_secret, ephemeral_spki) = ecdh_p256::compute_p256_shared_secret(public_key_point_hex)?;
    let enc_type = match cryptor {
        Cryptor::Aes256Gcm => ENC_AES256_GCM_P256,
        Cryptor::ChaCha20Poly1305 => ENC_CHACHA20_POLY1305_P256,
    };
    encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ephemeral_spki, enc_type, envelop)
}

fn encrypt_envelop_ecdh_p384(cryptor: Cryptor, key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let public_key_point_hex = &envelop.public_part;
    let (shared_secret, ephemeral_spki) = ecdh_p384::compute_p384_shared_secret(public_key_point_hex)?;
    let enc_type = match cryptor {
        Cryptor::Aes256Gcm => ENC_AES256_GCM_P384,
        Cryptor::ChaCha20Poly1305 => ENC_CHACHA20_POLY1305_P384,
    };
    encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ephemeral_spki, enc_type, envelop)
}

fn encrypt_envelop_ecdh_x25519(cryptor: Cryptor, key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let public_key_point_hex = &envelop.public_part;
    let (shared_secret, ephemeral_spki) = ecdh_x25519::compute_x25519_shared_secret(public_key_point_hex)?;
    let enc_type = match cryptor {
        Cryptor::Aes256Gcm => ENC_AES256_GCM_X25519,
        Cryptor::ChaCha20Poly1305 => ENC_CHACHA20_POLY1305_X25519,
    };
    encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ephemeral_spki, enc_type, envelop)
}

fn encrypt_envelop_ecdh_kyber1204(cryptor: Cryptor, key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let public_key_point_hex = &envelop.public_part;
    let (shared_secret, ephemeral_spki) = ecdh_kyber1024::compute_kyber1024_shared_secret(public_key_point_hex)?;
    let enc_type = match cryptor {
        Cryptor::Aes256Gcm => ENC_AES256_GCM_KYBER1204,
        Cryptor::ChaCha20Poly1305 => ENC_CHACHA20_POLY1305_KYBER1204,
    };
    encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ephemeral_spki, enc_type, envelop)
}

fn encrypt_envelop_ecdh_ml_kem(cryptor: Cryptor, key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let public_key_base64 = &envelop.public_part;
    let public_key = opt_result!(decode_base64(public_key_base64), "Decode ML-KEM public key from base64 failed: {}");
    let (shared_secret, ciphertext, ml_kem_algo) = util_mlkem::try_ml_kem_encapsulate(&public_key)?;
    let enc_type = match (cryptor, ml_kem_algo) {
        (Cryptor::Aes256Gcm, MlKemAlgo::MlKem768) => ENC_AES256_GCM_MLKEM768,
        (Cryptor::Aes256Gcm, MlKemAlgo::MlKem1024) => ENC_AES256_GCM_MLKEM1024,
        (Cryptor::ChaCha20Poly1305, MlKemAlgo::MlKem768) => ENC_CHACHA20_POLY1305_MLKEM768,
        (Cryptor::ChaCha20Poly1305, MlKemAlgo::MlKem1024) => ENC_CHACHA20_POLY1305_MLKEM1024,
    };
    encrypt_envelop_shared_secret(cryptor, key, &shared_secret, &ciphertext, enc_type, envelop)
}

fn encrypt_envelop_shared_secret(cryptor: Cryptor,
                                 key: &[u8],
                                 shared_secret: &[u8],
                                 ephemeral_spki: &[u8],
                                 enc_type: &str,
                                 envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let shared_key = util::simple_kdf(shared_secret);
    let nonce = util::make_nonce();
    let key_nonce = KeyNonce { k: &shared_key, n: nonce.as_ref() };

    let encrypted_key = crypto_simple::encrypt(
        cryptor, &key_nonce, key)?;

    let wrap_key = WrapKey {
        header: WrapKeyHeader::from(enc_type, ephemeral_spki),
        nonce: nonce.0.clone(),
        encrypted_data: encrypted_key,
    };
    let encoded_wrap_key = wrap_key.encode()?;

    Ok(TinyEncryptEnvelop {
        r#type: envelop.r#type,
        kid: envelop.kid.clone(),
        desc: None,
        encrypted_key: encoded_wrap_key,
    })
}

fn encrypt_envelop_rsa(key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let rsa_public_key = opt_result!(util_rsa::parse_spki(&envelop.public_part), "Parse RSA public key failed: {}");
    let mut rng = rand::thread_rng();
    let encrypted_key = opt_result!(rsa_public_key.encrypt(&mut rng, Pkcs1v15Encrypt, key), "RSA public key encrypt failed: {}");
    Ok(TinyEncryptEnvelop {
        r#type: envelop.r#type,
        kid: envelop.kid.clone(),
        desc: None,
        encrypted_key: util::encode_base64(&encrypted_key),
    })
}

fn encrypt_envelop_gpg(key: &[u8], envelop: &TinyEncryptConfigEnvelop) -> XResult<TinyEncryptEnvelop> {
    let encrypted_key = opt_result!(util_gpg::gpg_encrypt(&envelop.public_part, key), "GPG encrypt failed: {}");
    Ok(TinyEncryptEnvelop {
        r#type: envelop.r#type,
        kid: envelop.kid.clone(),
        desc: None,
        encrypted_key,
    })
}

fn get_compress_level(cmd_encrypt: &CmdEncrypt, path: &str, pem_output: bool) -> Option<u32> {
    let mut auto_compress = false;
    let path_parts = path.split(".").collect::<Vec<_>>();
    let path_ext = path_parts[path_parts.len() - 1].to_lowercase();
    if let Some(auto_compress_file_exts) = util_env::get_auto_compress_file_exts() {
        auto_compress = auto_compress_file_exts.contains(&path_ext);
        debugging!("File ext: {} matches auto compress exts: {:?}", path_ext, auto_compress_file_exts);
    }

    if auto_compress || cmd_encrypt.compress || util_env::get_default_compress().unwrap_or(false) {
        Some(cmd_encrypt.compress_level.unwrap_or_else(|| Compression::default().level()))
    } else if pem_output {
        Some(Compression::best().level())
    } else {
        None
    }
}