Skip to main content

magic_crypt/ciphers/
aes256.rs

1use alloc::vec::Vec;
2#[cfg(feature = "std")]
3use std::io::{Read, Write};
4
5#[cfg(feature = "std")]
6use aes::cipher::{
7    array::ArraySize,
8    typenum::{IsGreaterOrEqual, PartialDiv, True, U16},
9};
10use aes::{
11    Aes256,
12    cipher::{
13        BlockModeDecrypt, BlockModeEncrypt, Iv, Key, KeyIvInit, array::Array, block_padding::Pkcs7,
14    },
15};
16use md5::{Digest, Md5};
17use sha2::Sha256;
18
19#[cfg(feature = "std")]
20use crate::functions::{
21    decrypt_reader_to_writer as decrypt_stream, encrypt_reader_to_writer as encrypt_stream,
22    to_blocks,
23};
24use crate::{MagicCryptError, MagicCryptTrait};
25
26type Aes256CbcEnc = cbc::Encryptor<Aes256>;
27type Aes256CbcDec = cbc::Decryptor<Aes256>;
28
29#[cfg(feature = "std")]
30const BLOCK_SIZE: usize = 16;
31
32/// This struct can help you encrypt or decrypt data via AES-256 in a quick way.
33#[derive(Clone)]
34pub struct MagicCrypt256 {
35    key: Key<Aes256CbcEnc>,
36    iv:  Iv<Aes256CbcEnc>,
37}
38
39impl_cipher_secret_traits!(MagicCrypt256);
40
41impl MagicCryptTrait for MagicCrypt256 {
42    fn new<S: AsRef<[u8]>, V: AsRef<[u8]>>(key: S, iv: Option<V>) -> MagicCrypt256 {
43        let iv = match iv {
44            Some(s) => {
45                let mut hasher = Md5::new();
46                hasher.update(s.as_ref());
47
48                hasher.finalize()
49            },
50            None => Array::default(),
51        };
52
53        let key = {
54            let mut hasher = Sha256::new();
55            hasher.update(key.as_ref());
56
57            hasher.finalize()
58        };
59
60        MagicCrypt256 {
61            key,
62            iv,
63        }
64    }
65
66    #[inline]
67    fn encrypt_to_bytes<T: ?Sized + AsRef<[u8]>>(&self, data: &T) -> Vec<u8> {
68        let data = data.as_ref();
69
70        let cipher = Aes256CbcEnc::new(&self.key, &self.iv);
71
72        cipher.encrypt_padded_vec::<Pkcs7>(data)
73    }
74
75    #[cfg(feature = "std")]
76    fn encrypt_reader_to_bytes(&self, reader: &mut dyn Read) -> Result<Vec<u8>, MagicCryptError> {
77        use aes::cipher::block_padding::NoPadding;
78
79        let mut final_result = Vec::new();
80
81        let data_length = reader.read_to_end(&mut final_result)?;
82
83        let padding_length = BLOCK_SIZE - (data_length % BLOCK_SIZE);
84        let final_length = data_length + padding_length;
85
86        // PKCS7 padding requires that the padding bytes be equal to the number of padding bytes added.
87        let padding_byte = padding_length as u8;
88
89        final_result.resize(final_length, padding_byte);
90
91        let cipher = Aes256CbcEnc::new(&self.key, &self.iv);
92
93        cipher.encrypt_padded::<NoPadding>(&mut final_result, final_length).unwrap();
94
95        Ok(final_result)
96    }
97
98    #[cfg(feature = "std")]
99    fn encrypt_reader_to_writer2<
100        N: ArraySize + PartialDiv<U16> + IsGreaterOrEqual<U16, Output = True>,
101    >(
102        &self,
103        reader: &mut dyn Read,
104        writer: &mut dyn Write,
105    ) -> Result<(), MagicCryptError> {
106        let mut cipher = Aes256CbcEnc::new(&self.key, &self.iv);
107
108        encrypt_stream::<N>(reader, writer, BLOCK_SIZE, |buffer| {
109            cipher.encrypt_blocks(to_blocks(buffer));
110        })
111    }
112
113    #[inline]
114    fn decrypt_bytes_to_bytes<T: ?Sized + AsRef<[u8]>>(
115        &self,
116        bytes: &T,
117    ) -> Result<Vec<u8>, MagicCryptError> {
118        let bytes = bytes.as_ref();
119
120        let cipher = Aes256CbcDec::new(&self.key, &self.iv);
121
122        let final_result = cipher.decrypt_padded_vec::<Pkcs7>(bytes)?;
123
124        Ok(final_result)
125    }
126
127    #[cfg(feature = "std")]
128    fn decrypt_reader_to_bytes(&self, reader: &mut dyn Read) -> Result<Vec<u8>, MagicCryptError> {
129        let mut final_result = Vec::new();
130
131        reader.read_to_end(&mut final_result)?;
132
133        let cipher = Aes256CbcDec::new(&self.key, &self.iv);
134
135        let data_length = cipher.decrypt_padded::<Pkcs7>(&mut final_result)?.len();
136
137        final_result.truncate(data_length);
138
139        Ok(final_result)
140    }
141
142    #[cfg(feature = "std")]
143    fn decrypt_reader_to_writer2<
144        N: ArraySize + PartialDiv<U16> + IsGreaterOrEqual<U16, Output = True>,
145    >(
146        &self,
147        reader: &mut dyn Read,
148        writer: &mut dyn Write,
149    ) -> Result<(), MagicCryptError> {
150        let mut cipher = Aes256CbcDec::new(&self.key, &self.iv);
151
152        decrypt_stream::<N>(reader, writer, BLOCK_SIZE, |buffer| {
153            cipher.decrypt_blocks(to_blocks(buffer));
154        })
155    }
156}