use crate::crypto::DatCryptoAlgorithm::{IvAes128Gcm, IvAes256Gcm};
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::inout::InOutBuf;
use aes_gcm::{AeadInOut, AesGcm, Key, KeyInit, Nonce};
use std::fmt::Display;
use std::str::FromStr;
#[repr(u8)]
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum DatCryptoAlgorithm {
IvAes128Gcm,
IvAes256Gcm,
}
impl DatCryptoAlgorithm {
#[inline]
pub fn list() -> &'static [DatCryptoAlgorithm] {
&[IvAes128Gcm, IvAes256Gcm]
}
#[inline]
pub fn as_str(&self) -> &'static str {
match self {
IvAes128Gcm => "IV-AES128-GCM",
IvAes256Gcm => "IV-AES256-GCM",
}
}
}
const IV_LEN: usize = 12;
const TAG_LEN: usize = 16;
impl FromStr for DatCryptoAlgorithm {
type Err = DatError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"IV-AES128-GCM" => Ok(IvAes128Gcm),
"IV-AES256-GCM" => Ok(IvAes256Gcm),
_ => Err(DatError::UnknownCryptoAlgorithm),
}
}
}
impl Display for DatCryptoAlgorithm {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Clone)]
pub enum DatCrypto {
IvAes128Gcm(AesGcm<aes_gcm::aes::Aes128, U12>, Array<u8, U16>),
IvAes256Gcm(AesGcm<aes_gcm::aes::Aes256, U12>, Array<u8, U32>),
}
impl DatCrypto {
pub fn generate(algorithm: DatCryptoAlgorithm) -> Self {
match algorithm {
IvAes128Gcm => {
let key = Key::<aes_gcm::aes::Aes128>::generate();
let cipher = AesGcm::<aes_gcm::aes::Aes128, U12>::new(&key);
Self::IvAes128Gcm(cipher, key)
}
IvAes256Gcm => {
let key = Key::<aes_gcm::aes::Aes256>::generate();
let cipher = AesGcm::<aes_gcm::aes::Aes256, U12>::new(&key);
Self::IvAes256Gcm(cipher, key)
}
}
}
pub fn from_key(algorithm: DatCryptoAlgorithm, key: &[u8]) -> Result<Self, DatError> {
match algorithm {
IvAes128Gcm => {
let key = Key::<aes_gcm::aes::Aes128>::try_from(key)
.map_err(|_| DatError::InvalidCryptoKey)?;
let cipher = AesGcm::<aes_gcm::aes::Aes128, U12>::new(&key);
Ok(Self::IvAes128Gcm(cipher, key))
}
IvAes256Gcm => {
let key = Key::<aes_gcm::aes::Aes256>::try_from(key)
.map_err(|_| DatError::InvalidCryptoKey)?;
let cipher = AesGcm::<aes_gcm::aes::Aes256, U12>::new(&key);
Ok(Self::IvAes256Gcm(cipher, key))
}
}
}
#[inline]
pub fn algorithm(&self) -> DatCryptoAlgorithm {
match self {
Self::IvAes128Gcm(_, _) => IvAes128Gcm,
Self::IvAes256Gcm(_, _) => IvAes256Gcm,
}
}
pub fn export_key(&self) -> Box<[u8]> {
match self {
Self::IvAes128Gcm(_, key) => Box::from(key.as_slice()),
Self::IvAes256Gcm(_, key) => Box::from(key.as_slice()),
}
}
#[inline]
pub fn key_base64_len(&self) -> usize {
match self {
Self::IvAes128Gcm(_, _) => 22,
Self::IvAes256Gcm(_, _) => 43,
}
}
pub fn encrypt(&self, body: &[u8]) -> Result<Vec<u8>, DatError> {
if body.is_empty() {
return Ok(vec![]);
}
let iv: Array<u8, U12> = Nonce::generate();
let mut enc_data = Vec::with_capacity(IV_LEN + body.len() + TAG_LEN);
enc_data.extend_from_slice(&iv);
enc_data.extend_from_slice(body);
let inout = InOutBuf::from(&mut enc_data[IV_LEN..]);
let tag = match self {
Self::IvAes128Gcm(cipher, _) => cipher.encrypt_inout_detached(&iv, &[], inout),
Self::IvAes256Gcm(cipher, _) => cipher.encrypt_inout_detached(&iv, &[], inout),
}.map_err(|_| DatError::EncryptError)?;
enc_data.extend_from_slice(&tag.as_slice());
Ok(enc_data)
}
pub fn decrypt(&self, mut data: Vec<u8>) -> Result<Vec<u8>, DatError> {
if data.is_empty() {
return Ok(Vec::with_capacity(0));
}
if data.len() <= IV_LEN {
return Err(DatError::DecryptError);
}
let mut secure = data.split_off(IV_LEN);
match self {
Self::IvAes128Gcm(cipher, _) => {
AeadInOut::decrypt_in_place(cipher, data.as_slice().try_into().map_err(|_| DatError::DecryptError)?, &[], &mut secure)
},
Self::IvAes256Gcm(cipher, _) => {
AeadInOut::decrypt_in_place(cipher, data.as_slice().try_into().map_err(|_| DatError::DecryptError)?, &[], &mut secure)
},
}.map_err(|_| DatError::DecryptError)?;
Ok(secure)
}
}