enprot 0.5.19

// Copyright (c) 2018-2026 [Ribose Inc](https://www.ribose.com).
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

use std::collections::BTreeMap;

use crate::cipher;
use crate::cipher::{format_cipher_extfield, parse_cipher_extfield};
use crate::crypto::{self, CryptoPolicy};
use crate::error::{Error, Result};
use crate::etree;
use crate::pbkdf::{PBKDFCache, derive_key, parse_phc};

#[tracing::instrument(skip(pt, password, rng, pbkdfopts, cipheropts, cache, policy), fields(bytes = pt.len(), alg = %cipheropts.alg))]
pub fn encrypt(
    pt: Vec<u8>,
    password: &str,
    rng: &mut Option<botan::RandomNumberGenerator>,
    pbkdfopts: &etree::PBKDFOptions,
    cipheropts: &etree::CipherOptions,
    cache: &mut Option<PBKDFCache>,
    policy: &dyn CryptoPolicy,
) -> Result<(Vec<u8>, BTreeMap<String, String>)> {
    warn_legacy_pbkdf(&pbkdfopts.alg);

    policy
        .check_cipher_alg(&cipheropts.alg)
        .map_err(Error::Policy)?;

    let enc = cipher::encryption(&cipheropts.alg)?;
    let key_len = enc.key_len_max();
    let (master_key, pbkdf) = derive_key(password, key_len, rng, pbkdfopts, cache, policy)?;

    let mut extfields: BTreeMap<String, String> = BTreeMap::new();
    if let Some(p) = pbkdf {
        extfields.insert("pbkdf".to_string(), p);
    }

    let (key, iv) = compute_iv(cipheropts, &master_key, &pt, key_len, enc.nonce_len(), rng)?;

    if !cipheropts.alg.starts_with("aes-256-siv") {
        extfields.insert(
            "cipher".to_string(),
            format_cipher_extfield(&cipheropts.alg, &iv)?,
        );
    }

    policy
        .check_cipher(&cipheropts.alg, &key, &iv, &[])
        .map_err(Error::Policy)?;

    let pt_final = apply_compression(pt, cipheropts.compress, &mut extfields)?;

    let mut enc = enc;
    Ok((enc.process(&key, &iv, &[], &pt_final)?, extfields))
}

/// Warn when the deprecated legacy PBKDF is in use.
fn warn_legacy_pbkdf(alg: &str) {
    if alg == "legacy" {
        eprintln!(
            "Warning: --pbkdf legacy uses an unsalted SHA3-512 truncation and is \
             retained only for compatibility with old blobs; use argon2, scrypt, \
             or pbkdf2-sha{{256,512}} for new encryption."
        );
    }
}

/// Derive the AEAD key and nonce from the master key, plaintext, and
/// cipher options. Three modes:
///
/// - **Deterministic** (`-det` suffix): HKDF domain-separation into
///   enc-key + iv-key, then HMAC-derived nonce from plaintext. Same
///   `(password, plaintext)` → same `(key, iv)` → same ciphertext.
/// - **Random** (GCM etc.): RNG-generated or user-supplied IV.
/// - **SIV** (deterministic by design): no IV, master key used directly.
fn compute_iv(
    cipheropts: &etree::CipherOptions,
    master_key: &[u8],
    pt: &[u8],
    key_len: usize,
    nonce_len: usize,
    rng: &mut Option<botan::RandomNumberGenerator>,
) -> Result<(Vec<u8>, Vec<u8>)> {
    let is_det = cipheropts.alg.ends_with("-det");
    let needs_iv = !cipheropts.alg.starts_with("aes-256-siv");

    if is_det {
        if !needs_iv {
            return Err(Error::Cipher(format!(
                "{} does not support deterministic mode (SIV is already deterministic)",
                cipheropts.alg
            )));
        }
        let enc_key = crypto::hkdf_sha256(master_key, b"enprot-enc", key_len)?;
        let iv_key = crypto::hkdf_sha256(master_key, b"enprot-iv", 32)?;
        let iv_full = crypto::hmac_sha256(&iv_key, pt)?;
        let iv = iv_full[..nonce_len].to_vec();
        Ok((enc_key, iv))
    } else if needs_iv {
        let iv = match &cipheropts.iv {
            Some(v) => v.clone(),
            None => rng
                .as_mut()
                .ok_or(Error::InvalidArg {
                    arg: "rng",
                    reason: "Missing RNG for non-deterministic encrypt".to_string(),
                })?
                .read(nonce_len)
                .map_err(Error::botan)?,
        };
        Ok((master_key.to_vec(), iv))
    } else if cipheropts.iv.is_some() {
        Err(Error::Cipher("IV was supplied but not expected".into()))
    } else {
        Ok((master_key.to_vec(), Vec::new()))
    }
}

/// Optionally compress plaintext before encryption. Adds the
/// `compress:zlib` extfield only when compression actually reduced
/// the size. Returns the original bytes unchanged when compression
/// is disabled or ineffective.
fn apply_compression(
    pt: Vec<u8>,
    compress: bool,
    extfields: &mut BTreeMap<String, String>,
) -> Result<Vec<u8>> {
    if !compress {
        return Ok(pt);
    }
    let (compressed, did_compress) = crate::compress::compress(&pt)?;
    if did_compress {
        extfields.insert(
            "compress".to_string(),
            crate::compress::COMPRESS_EXTFIELD.to_string(),
        );
        Ok(compressed)
    } else {
        Ok(pt)
    }
}

#[tracing::instrument(skip(ct, password, pbkdf, cipher, cache, policy), fields(bytes = ct.len()))]
pub fn decrypt(
    ct: Vec<u8>,
    password: &str,
    pbkdf: &Option<&String>,
    cipher: &Option<&String>,
    compress: &Option<&String>,
    cache: &mut Option<PBKDFCache>,
    policy: &dyn CryptoPolicy,
) -> Result<Vec<u8>> {
    let (cipher_alg, iv) = match cipher {
        Some(s) => parse_cipher_extfield(s)?,
        None => (cipher::DEFAULT_CIPHER_ALG.to_string(), Vec::new()),
    };

    policy
        .check_cipher_alg(&cipher_alg)
        .map_err(Error::Policy)?;

    let dec = cipher::decryption(&cipher_alg)?;
    let key_len = dec.key_len_max();
    let master_key = derive_decrypt_key(password, pbkdf, key_len, cache, policy)?;
    let key = recover_key(&cipher_alg, master_key, key_len)?;

    policy
        .check_cipher(&cipher_alg, &key, &iv, &[])
        .map_err(Error::Policy)?;

    let mut dec = dec;
    let pt = dec.process(&key, &iv, &[], &ct)?;

    if let Some(alg) = compress
        && alg.as_str() == crate::compress::COMPRESS_EXTFIELD
    {
        crate::compress::decompress(&pt)
    } else {
        Ok(pt)
    }
}

/// Derive the master decryption key from the password + PHC extfield.
/// Falls back to the legacy unsalted SHA3-512 path when no PHC string
/// is present (old blobs encrypted before PBKDF extfields existed).
fn derive_decrypt_key(
    password: &str,
    pbkdf: &Option<&String>,
    key_len: usize,
    cache: &mut Option<PBKDFCache>,
    policy: &dyn CryptoPolicy,
) -> Result<Vec<u8>> {
    let mut no_rng: Option<botan::RandomNumberGenerator> = None;
    let pbkdfopts = match pbkdf {
        Some(p) => {
            let (alg, params, salt) = parse_phc(p)?;
            etree::PBKDFOptions {
                alg,
                saltlen: 0,
                salt: Some(salt),
                msec: None,
                params: Some(params),
            }
        }
        None => etree::PBKDFOptions {
            alg: "legacy".to_string(),
            saltlen: 0,
            salt: None,
            msec: None,
            params: None,
        },
    };
    Ok(derive_key(password, key_len, &mut no_rng, &pbkdfopts, cache, policy)?.0)
}

/// Recover the AEAD key from the master key. For deterministic
/// variants (`-det` suffix), the encrypt path domain-separated the
/// key via HKDF; decrypt must do the same to recover it. For all
/// other algorithms, the master key is used directly.
fn recover_key(cipher_alg: &str, master_key: Vec<u8>, key_len: usize) -> Result<Vec<u8>> {
    if cipher_alg.ends_with("-det") {
        crypto::hkdf_sha256(&master_key, b"enprot-enc", key_len)
    } else {
        Ok(master_key)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn cipher_extfield_roundtrip() {
        // aes-256-gcm with IV
        let s = "aes-256-gcm$iv=MDEyMzQ1Njc4OWFiY2RlZg==";
        let (alg, iv) = parse_cipher_extfield(s).unwrap();
        assert_eq!(alg, "aes-256-gcm");
        assert_eq!(iv, b"0123456789abcdef");
    }

    #[test]
    fn cipher_extfield_default_when_unset() {
        // When the cipher: field is absent on a legacy ENCRYPTED block,
        // prot::decrypt falls back to cipher::DEFAULT_CIPHER_ALG.
        let (alg, iv) = (
            crate::cipher::DEFAULT_CIPHER_ALG.to_string(),
            Vec::<u8>::new(),
        );
        assert_eq!(alg, "aes-256-siv");
        assert!(iv.is_empty());
    }

    #[test]
    fn phc_extfield_roundtrip() {
        let s = "$argon2$m=65536,p=4,t=3$MDEyMzQ1Njc4OWFiY2RlZg==";
        let (alg, params, salt) = parse_phc(s).unwrap();
        assert_eq!(alg, "argon2");
        assert_eq!(params.get("m"), Some(&65536));
        assert_eq!(params.get("t"), Some(&3));
        assert_eq!(params.get("p"), Some(&4));
        assert_eq!(salt, b"0123456789abcdef");
    }
}