byte_aes/
cryptor.rs

1use aes::cipher::{
2    crypto_common::generic_array::GenericArray, typenum::U32, BlockDecrypt, BlockEncrypt,
3};
4use aes::{cipher::KeyInit, Aes256};
5
6impl TryFrom<&str> for Aes256Cryptor {
7    type Error = std::io::Error;
8
9    fn try_from(value: &str) -> Result<Self, Self::Error> {
10        if value.len() != 32 {
11            return Err(std::io::Error::new(
12                std::io::ErrorKind::InvalidData,
13                "The number of bytes of the key shall be 32",
14            ));
15        }
16        let mut key = [0u8; 32];
17        key.copy_from_slice(value.as_bytes());
18        Ok(Aes256Cryptor::new(key))
19    }
20}
21
22impl TryFrom<String> for Aes256Cryptor {
23    type Error = std::io::Error;
24
25    fn try_from(value: String) -> Result<Self, Self::Error> {
26        Aes256Cryptor::try_from(&value as &str)
27    }
28}
29
30impl TryFrom<&String> for Aes256Cryptor {
31    type Error = std::io::Error;
32
33    fn try_from(value: &String) -> Result<Self, Self::Error> {
34        Aes256Cryptor::try_from(&value as &str)
35    }
36}
37
38#[derive(Clone, Debug)]
39pub struct Aes256Cryptor {
40    key: [u8; 32],
41    aes_256: Aes256,
42}
43
44impl Aes256Cryptor {
45    pub fn new(key: [u8; 32]) -> Self {
46        Self {
47            // Here U32 is not same as u32. U32 creates our GenericArray Key which is 32 byte long. If the key length is not 32 bytes in length, then this from_slice() call fails with assertion error
48            key,
49            aes_256: Aes256::new(GenericArray::<u8, U32>::from_slice(&key)),
50        }
51    }
52
53    pub fn key(&self) -> &[u8] {
54        &self.key
55    }
56
57    pub fn encrypt<T, U>(&self, plaintext: T) -> Vec<u8>
58    where
59        BytesWrapper<BytesWrapper<T>>: IntoBytes<U>,
60    {
61        let mut blocks =
62            super::get_generic_array(BytesWrapper(BytesWrapper(plaintext)).into_bytes(), false);
63        // Encrypt the single unit at once (This Unit will contain all the blocks)
64        self.aes_256.encrypt_blocks(blocks.as_mut_slice());
65        blocks.concat().into_iter().collect()
66    }
67
68    pub fn decrypt<T, U>(&self, ciphertext: T) -> std::io::Result<Vec<u8>>
69    where
70        BytesWrapper<BytesWrapper<T>>: IntoBytes<U>,
71    {
72        let raw_bytes = BytesWrapper(BytesWrapper(ciphertext)).into_bytes();
73        if raw_bytes.len() % 16 != 0 {
74            return Err(std::io::Error::new(
75                std::io::ErrorKind::InvalidData,
76                "The number of bytes of the encrypted data shall be multiple of 16",
77            ));
78        }
79        let mut blocks = super::get_generic_array(raw_bytes, true);
80        // Decrypt the single unit at once (This Unit will contain all the blocks)
81        self.aes_256.decrypt_blocks(blocks.as_mut_slice());
82        // Concat the decrypted block which the deref_generic_block hold and turn then in a Vec<u8>
83        let decrypted_bytes = blocks.concat().into_iter().collect::<Vec<u8>>();
84
85        if let Some(v) = decrypted_bytes.last() {
86            if *v == 16 && decrypted_bytes[decrypted_bytes.len() - 16..] == [16u8; 16] {
87                return Ok(decrypted_bytes[..decrypted_bytes.len() - 16].to_vec());
88            } else if *v < 16 {
89                return Ok(decrypted_bytes[..decrypted_bytes.len() - (*v as usize)].to_vec());
90            }
91            Err(std::io::Error::new(
92                std::io::ErrorKind::InvalidData,
93                format!(
94                    "Invalid encrypted data, the padding number cannot be {}",
95                    *v
96                ),
97            ))
98        } else {
99            Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "The number of bytes of the encrypted data shall be at least 16 even if the original data is empty"))
100        }
101    }
102}
103
104pub struct BytesWrapper<T>(T);
105
106pub trait IntoBytes<T> {
107    fn into_bytes(self) -> Vec<u8>;
108}
109
110impl<T: Into<Vec<u8>>> IntoBytes<T> for T {
111    fn into_bytes(self) -> Vec<u8> {
112        self.into()
113    }
114}
115
116impl<T: IntoBytes<U>, U> IntoBytes<U> for BytesWrapper<T> {
117    fn into_bytes(self) -> Vec<u8> {
118        self.0.into_bytes()
119    }
120}
121
122impl IntoBytes<()> for BytesWrapper<BytesWrapper<&String>> {
123    fn into_bytes(self) -> Vec<u8> {
124        self.0 .0.to_owned().into_bytes()
125    }
126}