use crate::error::DatError;
use aes_gcm::aead::array::Array;
use aes_gcm::aead::common::Generate;
use aes_gcm::aead::consts::{U12, U16, U32};
use aes_gcm::aead::Aead;
use aes_gcm::{Aes128Gcm, AesGcm};
use aes_gcm::{Aes256Gcm, Key, KeyInit, Nonce};
use crate::crypto_algorithm::CryptoAlgorithm;
pub enum CryptoKey {
AES128GCMN(Array<u8, U16>, AesGcm<aes_gcm::aes::Aes128, U12>),
AES256GCMN(Array<u8, U32>, AesGcm<aes_gcm::aes::Aes256, U12>),
}
impl CryptoKey {
pub fn algorithm(&self) -> CryptoAlgorithm {
match self {
CryptoKey::AES128GCMN(_,_) => CryptoAlgorithm::AES128GCMN,
CryptoKey::AES256GCMN(_,_) => CryptoAlgorithm::AES256GCMN,
}
}
pub fn from_bytes(algorithm: CryptoAlgorithm, data: &[u8]) -> Result<CryptoKey, DatError> {
Ok(match algorithm {
CryptoAlgorithm::AES128GCMN => {
let bytes = Array::<u8, U16>::try_from(data).map_err(|_| DatError::ParseCryptoKeyError)?;
CryptoKey::AES128GCMN(bytes, Aes128Gcm::new(&bytes))
},
CryptoAlgorithm::AES256GCMN => {
let bytes = Array::<u8, U32>::try_from(data).map_err(|_| DatError::ParseCryptoKeyError)?;
CryptoKey::AES256GCMN(bytes, Aes256Gcm::new(&bytes))
},
})
}
pub fn to_bytes(&self) -> Box<[u8]> {
match self {
CryptoKey::AES128GCMN(c, _) => Box::from(c.as_slice()),
CryptoKey::AES256GCMN(c, _) => Box::from(c.as_slice()),
}
}
pub fn generate(algorithm: CryptoAlgorithm) -> Self {
match algorithm {
CryptoAlgorithm::AES128GCMN => {
let bytes = Key::<Aes128Gcm>::generate();
CryptoKey::AES128GCMN(bytes, Aes128Gcm::new(&bytes))
},
CryptoAlgorithm::AES256GCMN => {
let bytes = Key::<Aes256Gcm>::generate();
CryptoKey::AES256GCMN(bytes, Aes256Gcm::new(&bytes))
},
}
}
pub fn encrypt(&self, body: &[u8]) -> Result<Vec<u8>, DatError> {
if body.is_empty() {
return Ok(vec![]);
}
let nonce = Nonce::generate();
let mut enc_data = Vec::with_capacity(50 + nonce.len() + body.len());
enc_data.extend_from_slice(&nonce);
match self {
CryptoKey::AES128GCMN(_, c) => {
enc_data.extend_from_slice(&c.encrypt(&nonce, body).map_err(|_| DatError::EncryptError)?);
},
CryptoKey::AES256GCMN(_, c) => {
enc_data.extend_from_slice(&c.encrypt(&nonce, body).map_err(|_| DatError::EncryptError)?);
},
}
Ok(enc_data)
}
pub fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, DatError> {
if data.is_empty() {
return Ok(vec![]);
}
let nonce = data[..12].try_into().map_err(|_| DatError::DecryptError)?;
match self {
CryptoKey::AES128GCMN(_, c) => {
c.decrypt(&nonce, &data[12..]).map_err(|_| DatError::DecryptError)
},
CryptoKey::AES256GCMN(_, c) => {
c.decrypt(&nonce, &data[12..]).map_err(|_| DatError::DecryptError)
},
}
}
}
impl Clone for CryptoKey {
fn clone(&self) -> Self {
match self {
CryptoKey::AES128GCMN(b, c) => CryptoKey::AES128GCMN(b.clone(), c.clone()),
CryptoKey::AES256GCMN(b, c) => CryptoKey::AES256GCMN(b.clone(), c.clone()),
}
}
}