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};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CryptAlgorithm {
AES128GCMN,
AES256GCMN,
}
impl CryptAlgorithm {
pub fn to_str(&self) -> &'static str {
match self {
CryptAlgorithm::AES128GCMN => "AES128GCMN",
CryptAlgorithm::AES256GCMN => "AES256GCMN",
}
}
pub fn from_str(s: &str) -> Result<CryptAlgorithm, DatError> {
match s {
"AES128GCMN" => Ok(CryptAlgorithm::AES128GCMN),
"AES256GCMN" => Ok(CryptAlgorithm::AES256GCMN),
_ => Err(DatError::CryptError(format!("UnknownCryptAlgorithm: {}", s))),
}
}
}
pub enum CryptKey {
AES128GCMN(Array<u8, U16>),
AES256GCMN(Array<u8, U32>),
}
pub enum Cipher {
AES128GCMN(AesGcm<aes_gcm::aes::Aes128, U12>),
AES256GCMN(AesGcm<aes_gcm::aes::Aes256, U12>),
}
impl CryptKey {
pub fn to_alg(&self) -> CryptAlgorithm {
match self {
CryptKey::AES128GCMN(_) => CryptAlgorithm::AES128GCMN,
CryptKey::AES256GCMN(_) => CryptAlgorithm::AES256GCMN,
}
}
pub fn from_bytes(alg: CryptAlgorithm, data: &[u8]) -> Result<CryptKey, DatError> {
Ok(match alg {
CryptAlgorithm::AES128GCMN => CryptKey::AES128GCMN(
Array::<u8, U16>::try_from(data)
.map_err(|e| DatError::CryptError(format!("invalid crypt key: {}: {}", alg.to_str(), e)))?
),
CryptAlgorithm::AES256GCMN => CryptKey::AES256GCMN(
Array::<u8, U32>::try_from(data)
.map_err(|e| DatError::CryptError(format!("invalid crypt key: {}: {}", alg.to_str(), e)))?
),
})
}
pub fn to_bytes(&self) -> Box<[u8]> {
match self {
CryptKey::AES128GCMN(c) => Box::from(c.as_slice()),
CryptKey::AES256GCMN(c) => Box::from(c.as_slice()),
}
}
pub fn to_cipher(&self) -> Cipher {
match self {
CryptKey::AES128GCMN(c) => Cipher::AES128GCMN(Aes128Gcm::new(c)),
CryptKey::AES256GCMN(c) => Cipher::AES256GCMN(Aes256Gcm::new(c)),
}
}
pub fn generate(alg: CryptAlgorithm) -> Self {
match alg {
CryptAlgorithm::AES128GCMN => CryptKey::AES128GCMN(Key::<Aes128Gcm>::generate()),
CryptAlgorithm::AES256GCMN => CryptKey::AES256GCMN(Key::<Aes256Gcm>::generate()),
}
}
}
impl Cipher {
pub fn encrypt(&self, body: &[u8]) -> Result<Vec<u8>, DatError> {
if body.is_empty() {
return Ok(vec![]);
}
macro_rules! aes_gcmn_encrypt {
($cipher:expr) => {{
let nonce = Nonce::generate();
let mut enc_data = nonce.to_vec();
enc_data.extend(
$cipher.encrypt(&nonce, body)
.map_err(|e| DatError::CryptError(format!("encrypt error {}", e)))?
);
Ok(enc_data)
}};
}
match self {
Cipher::AES128GCMN(c) => aes_gcmn_encrypt!(c),
Cipher::AES256GCMN(c) => aes_gcmn_encrypt!(c),
}
}
pub fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, DatError> {
if data.is_empty() {
return Ok(vec![]);
}
macro_rules! aes_gcmn_decrypt {
($cipher:expr) => {{
let nonce = data[..12].try_into()
.map_err(|_| DatError::CryptError(format!("decrypt nonce error")))?;
$cipher.decrypt(&nonce, &data[12..])
.map_err(|e| DatError::CryptError(format!("decrypt error {}", e)))
}};
}
match self {
Cipher::AES128GCMN(c) => aes_gcmn_decrypt!(c),
Cipher::AES256GCMN(c) => aes_gcmn_decrypt!(c),
}
}
}