use super::{Result, SDLanError};
use aes::cipher::{block_padding::Pkcs7, BlockDecryptMut, BlockEncryptMut, KeyIvInit};

type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;

pub fn aes_encrypt(key: &[u8], plain: &[u8]) -> Result<Vec<u8>> {
    let mut buf = Vec::new();
    buf.resize(plain.len() + 16, 0);

    let iv = &key[..16];
    match Aes256CbcEnc::new(key.into(), iv.into()).encrypt_padded_b2b_mut::<Pkcs7>(plain, &mut buf)
    {
        Err(e) => Err(SDLanError::EncryptError(format!("aes encrypt: {}", e))),
        Ok(v) => Ok(Vec::from(v)),
    }
}

pub fn aes_decrypt(key: &[u8], cipherd: &[u8]) -> Result<Vec<u8>> {
    let mut buf = Vec::new();
    let iv = &key[..16];
    buf.resize(cipherd.len() + 16, 0);
    match Aes256CbcDec::new(key.into(), iv.into())
        .decrypt_padded_b2b_mut::<Pkcs7>(cipherd, &mut buf)
    {
        Err(e) => Err(SDLanError::EncryptError(format!("aes decrypt: {}", e))),
        Ok(v) => Ok(Vec::from(v)),
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use rand::{thread_rng, Rng};

    fn gen_aes_key() -> [u8; 32] {
        let mut res = [0; 32];
        for i in 0..32 {
            let mut temp = thread_rng().gen_range(0..63) as u8;
            if temp <= 10 {
                temp += 0x30;
            } else if temp <= 36 {
                temp += 0x61;
            } else {
                temp += 0x41;
            }
            res[i as usize] = temp;
        }

        return res;
    }

    #[test]
    fn test_aes() -> Result<()> {
        use crate::utils::test_utils::generate_info;
        let key = gen_aes_key();
        println!("aes key: {:?}", key);
        let zero_msg = vec![];
        let short_msg = generate_info(10);
        let middle_msg = generate_info(1024);
        let big_msg = generate_info(65535);

        let msgs = [zero_msg, short_msg, middle_msg, big_msg];

        for msg in msgs.iter() {
            println!("test aes decrypt with {} bytes", msg.len());
            let encrypted = aes_encrypt(&key, msg.as_slice())?;
            let decrypted = aes_decrypt(&key, encrypted.as_slice())?;
            assert_eq!(&decrypted, msg);
            println!("test aes decrypt {} ok", msg.len());
        }

        Ok(())
    }
}