mycommon-utils 0.1.2

use base64::{Engine};
use base64::engine::general_purpose::STANDARD;
use crypto::{aes, symmetriccipher};
use crypto::aes::KeySize::KeySize256;
use crypto::blockmodes::PkcsPadding;
use crypto::buffer::{BufferResult, ReadBuffer, RefReadBuffer, RefWriteBuffer, WriteBuffer};
use rand::random;

/// Encrypt a buffer with the given key and iv using AES256/CBC/Pkcs encryption.
/// 加密(data:加密数据；key：密钥（长度为32的字符串）；iv：偏移量（长度为16的字符串）)
pub fn aes256_cbc_encrypt(
    data: &[u8],
    key: &str,
) -> String {
    let key = string_to_fixed_array_32(key);
    let mut iv = random::<[u8; 16]>();  // 生成一个 IV（初始化向量）
    // println!("iv:{:?}", iv);
    let output = encrypt(data, &key, &mut iv);
    // 将 IV 和加密数据合并
    let mut result = iv.to_vec();
    result.extend(&output.expect("Failed to encrypt data"));
    STANDARD.encode(&result) // 编码为 Base64 字符串
}

// Encrypt a buffer with the given key and iv using
// AES-256/CBC/Pkcs encryption.
fn encrypt(data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>, symmetriccipher::SymmetricCipherError> {
    // Create an encryptor instance of the best performing
    // type available for the platform.
    let mut encryptor = aes::cbc_encryptor(KeySize256, key, iv, PkcsPadding);
    let mut final_result = Vec::<u8>::new();
    let mut read_buffer = RefReadBuffer::new(data);
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);
    loop {
        let result = encryptor.encrypt(&mut read_buffer, &mut write_buffer, true)?;
        final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));
        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => { }
        }
    }
    Ok(final_result)
}

/// Decrypt a buffer with the given key and iv using AES256/CBC/Pkcs encryption.
/// 解密(data:加密数据；key：密钥（长度为32的字符串）；iv：偏移量（长度为16的字符串）)
pub fn aes256_cbc_decrypt(
    data: &str,
    key: &str
) -> String {
    // 解码 Base64 数据
    let data = STANDARD.decode(data).expect("Failed to decode Base64");
    // println!("Decoded data: {:?}", data);
    let key = string_to_fixed_array_32(key);
    // 解析 IV 和加密数据
    let (iv, data) = data.split_at(16);
    // println!("IV: {:?}", iv);
    // println!("Encrypted data: {:?}", data);
    let output = decrypt(data, &key, iv).expect("Failed to decrypt data");
    // 解密后转换为 UTF-8 字符串
    String::from_utf8(output).expect("Failed to convert to string")
}

fn decrypt(encrypted_data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>, symmetriccipher::SymmetricCipherError> {
    let mut decryptor = aes::cbc_decryptor(KeySize256, key, iv, PkcsPadding);
    let mut final_result = Vec::<u8>::new();
    let mut read_buffer = RefReadBuffer::new(encrypted_data);
    let mut buffer = [0; 4096];
    let mut write_buffer = RefWriteBuffer::new(&mut buffer);
    loop {
        let result = decryptor.decrypt(&mut read_buffer, &mut write_buffer, true)?;
        final_result.extend(write_buffer.take_read_buffer().take_remaining().iter().map(|&i| i));
        match result {
            BufferResult::BufferUnderflow => break,
            BufferResult::BufferOverflow => { }
        }
    }
    Ok(final_result)
}

/// 将字符串转为[u8; 32]
fn string_to_fixed_array_32(s: &str) -> [u8; 32] {
    s.as_bytes().try_into().expect("字符串转为[u8; 32]失败")
}