use aes::cipher::{BlockModeDecrypt, BlockModeEncrypt, KeyIvInit, block_padding::Pkcs7};
use hmac::{KeyInit, Mac};
use hybrid_array::Array;
use subtle::ConstantTimeEq;
use typenum::U32;
use crate::{
error::Result,
util::{PBKDF_SHA256_HMAC_OUT_SIZE, PbkdfSha256Hmac},
};
#[derive(Debug)]
pub(crate) struct DecryptError {}
pub(crate) fn decrypt_aes256(
iv: &[u8; 16],
data: Vec<u8>,
key: &Array<u8, U32>,
) -> Result<Vec<u8>, DecryptError> {
let mut data = data;
let decrypted_key_slice = cbc::Decryptor::<aes::Aes256>::new(key, iv.into())
.decrypt_padded::<Pkcs7>(&mut data)
.map_err(|_| DecryptError {})?;
let decrypted_len = decrypted_key_slice.len();
data.truncate(decrypted_len);
Ok(data)
}
pub(crate) fn decrypt_aes256_hmac(
iv: &[u8; 16],
mac: &[u8; 32],
data: Vec<u8>,
mac_key: &Array<u8, U32>,
key: &Array<u8, U32>,
) -> Result<Vec<u8>, DecryptError> {
let res = generate_mac(mac_key, iv, &data);
if res.ct_ne(mac).into() {
return Err(DecryptError {});
}
decrypt_aes256(iv, data, key)
}
pub(crate) fn encrypt_aes256_hmac(
data_dec: &[u8],
mac_key: &Array<u8, U32>,
key: &Array<u8, U32>,
) -> Result<([u8; 16], [u8; 32], Vec<u8>)> {
let rng = rand::rng();
let (iv, data) = encrypt_aes256_internal(rng, data_dec, key);
let mac = generate_mac(mac_key, &iv, &data);
Ok((iv, mac, data))
}
fn encrypt_aes256_internal(
mut rng: impl rand::Rng,
data_dec: &[u8],
key: &Array<u8, U32>,
) -> ([u8; 16], Vec<u8>) {
let mut iv = [0u8; 16];
rng.fill_bytes(&mut iv);
let data =
cbc::Encryptor::<aes::Aes256>::new(key, &iv.into()).encrypt_padded_vec::<Pkcs7>(data_dec);
(iv, data)
}
fn generate_mac(mac_key: &[u8], iv: &[u8], data: &[u8]) -> [u8; 32] {
let mut hmac =
PbkdfSha256Hmac::new_from_slice(mac_key).expect("hmac new_from_slice should not fail");
hmac.update(iv);
hmac.update(data);
let mac: [u8; PBKDF_SHA256_HMAC_OUT_SIZE] = (*hmac.finalize().into_bytes())
.try_into()
.expect("HMAC output size to be correct");
mac
}
#[cfg(test)]
mod tests {
use bitwarden_encoding::B64;
use hybrid_array::ArraySize;
use rand::SeedableRng;
use super::*;
fn generate_array<N: ArraySize>(offset: u8, increment: u8) -> Array<u8, N> {
Array::from_fn(|i| offset + i as u8 * increment)
}
fn generate_vec(length: usize, offset: u8, increment: u8) -> Vec<u8> {
(0..length).map(|i| offset + i as u8 * increment).collect()
}
#[test]
fn test_encrypt_aes256_internal() {
let key = generate_array(0, 1);
let rng = rand_chacha::ChaCha8Rng::from_seed([0u8; 32]);
let result = encrypt_aes256_internal(rng, "EncryptMe!".as_bytes(), &key);
assert_eq!(
result,
(
[
62, 0, 239, 47, 137, 95, 64, 214, 127, 91, 184, 232, 31, 9, 165, 161
],
vec![
214, 76, 187, 97, 58, 146, 212, 140, 95, 164, 177, 204, 179, 133, 172, 148
]
)
);
}
#[test]
fn test_generate_mac() {
let mac_key = generate_vec(16, 0, 16);
let iv = generate_vec(16, 0, 16);
let data = generate_vec(16, 0, 16);
let mac = generate_mac(&mac_key, &iv, &data);
assert!(mac.iter().any(|&b| b != 0));
}
#[test]
fn test_decrypt_aes256() {
let iv = generate_vec(16, 0, 1);
let iv: &[u8; 16] = iv.as_slice().try_into().unwrap();
let key = generate_array(0, 1);
let data: B64 = ("ByUF8vhyX4ddU9gcooznwA==").parse().unwrap();
let decrypted = decrypt_aes256(iv, data.into(), &key).unwrap();
assert_eq!(String::from_utf8(decrypted).unwrap(), "EncryptMe!");
}
}