1use crate::aes::{AESError, AES, BLOCK_SIZE};
7use crate::pkcs7;
8
9use std::error::Error;
10use std::fmt;
11use std::fmt::Display;
12
13pub struct Cipher {
39 aes: AES,
40 pad: pkcs7::Pad,
41 iv: [u8; BLOCK_SIZE],
42}
43
44impl Cipher {
45 pub fn new(iv: [u8; BLOCK_SIZE]) -> Self {
49 Cipher {
50 aes: AES::new(),
51 pad: pkcs7::Pad::new(BLOCK_SIZE),
52 iv,
53 }
54 }
55
56 pub fn init(&mut self, key: &[u8]) -> Result<(), AESError> {
60 self.aes.init(key)?;
61 Ok(())
62 }
63
64 pub fn write(&mut self, b: &[u8]) -> Result<Vec<u8>, AESError> {
66 let blocks = match self.pad.write(b) {
67 None => {
68 return Ok(vec![]);
69 }
70 Some(block) => block,
71 };
72
73 let mut result: Vec<u8> = Vec::with_capacity(blocks.len());
74 for block in blocks.chunks(BLOCK_SIZE) {
75 let mut ciphertext = self.encrypt_block(block)?;
76 result.append(&mut ciphertext);
77 }
78 Ok(result)
79 }
80
81 pub fn flush(&mut self) -> Result<Vec<u8>, AESError> {
83 let block = self.pad.flush();
84 Ok(self.encrypt_block(&block)?)
85 }
86
87 fn encrypt_block(&mut self, b: &[u8]) -> Result<Vec<u8>, AESError> {
88 let mut block = [0; BLOCK_SIZE];
89 block.copy_from_slice(&b);
90
91 for (elem, iv_elem) in block.iter_mut().zip(self.iv.iter()) {
92 *elem ^= iv_elem;
93 }
94
95 match self.aes.encrypt(&block) {
96 Ok(new_iv) => {
97 self.iv = new_iv;
98 Ok(new_iv.to_vec())
99 }
100 Err(err) => Err(err),
101 }
102 }
103}
104
105pub struct Decipher {
134 aes: AES,
135 unpad: pkcs7::Unpad,
136 iv: [u8; BLOCK_SIZE],
137 buffer: pkcs7::Pad,
138}
139
140#[derive(Debug, PartialEq)]
142pub enum DecipherError {
143 AES(AESError),
145 PKCS7(pkcs7::UnpadError),
147}
148
149impl Display for DecipherError {
150 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
151 write!(f, "DecipherError: {}", self.description())
152 }
153}
154
155impl Error for DecipherError {
156 fn description(&self) -> &str {
157 match self {
159 DecipherError::AES(err) => err.description(),
160 DecipherError::PKCS7(err) => err.description(),
161 }
162 }
163}
164
165impl Decipher {
166 pub fn new(iv: [u8; BLOCK_SIZE]) -> Self {
170 Decipher {
171 aes: AES::new(),
172 unpad: pkcs7::Unpad::new(BLOCK_SIZE),
173 iv,
174 buffer: pkcs7::Pad::new(BLOCK_SIZE),
175 }
176 }
177
178 pub fn init(&mut self, key: &[u8]) -> Result<(), AESError> {
182 self.aes.init(key)?;
183 Ok(())
184 }
185
186 pub fn write(&mut self, b: &[u8]) -> Result<Vec<u8>, DecipherError> {
188 let blocks = match self.buffer.write(b) {
189 Some(blocks) => blocks,
190 None => {
191 return Ok(vec![]);
192 }
193 };
194
195 let mut padded_result = Vec::with_capacity(blocks.len());
196 for ciphertext in blocks.chunks(BLOCK_SIZE) {
197 let mut block = [0; BLOCK_SIZE];
198 block.copy_from_slice(&ciphertext);
199
200 let mut cleartext = match self.aes.decrypt(&block) {
201 Err(err) => return Err(DecipherError::AES(err)),
202 Ok(out) => out,
203 };
204
205 for (elem, iv_elem) in cleartext.iter_mut().zip(self.iv.iter()) {
206 *elem ^= iv_elem;
207 }
208
209 self.iv = block;
210 padded_result.append(&mut cleartext.to_vec());
211 }
212
213 Ok(self.unpad.write(&padded_result).unwrap_or_default())
214 }
215
216 pub fn flush(&mut self) -> Result<Vec<u8>, DecipherError> {
218 match self.unpad.flush() {
219 Err(err) => Err(DecipherError::PKCS7(err)),
220 Ok(block) => Ok(block),
221 }
222 }
223}
224
225#[cfg(test)]
226mod tests {
227 use super::*;
228 use crate::common::hex_to_vec;
229
230 fn check(key: &str, iv: &str, cleartext: &str, ciphertext: &str) {
231 let key_vec = hex_to_vec(key);
232 let mut iv_arr = [0; BLOCK_SIZE];
233 iv_arr.copy_from_slice(&hex_to_vec(iv));
234
235 let mut cipher = Cipher::new(iv_arr);
236 cipher.init(&key_vec).expect("cipher init to not fail");
237
238 let mut actual = cipher
239 .write(cleartext.as_bytes())
240 .expect("cipher write to not fail");
241 actual.append(&mut cipher.flush().expect("cipher flush to not fail"));
242
243 assert_eq!(actual, hex_to_vec(ciphertext));
244
245 let mut decipher = Decipher::new(iv_arr);
246 decipher.init(&key_vec).expect("decipher init to not fail");
247 let mut back = decipher
248 .write(&hex_to_vec(ciphertext))
249 .expect("decipher write to not fail");
250 back.append(&mut decipher.flush().expect("decipher flush to not fail"));
251
252 assert_eq!(back, cleartext.as_bytes());
253 }
254
255 #[test]
256 fn it_should_not_fail_on_vec_0() {
257 check(
258 "00000000000000000000000000000000",
259 "00000000000000000000000000000000",
260 "hello",
261 "9834ed518cbc8fbe9af3c6ecb75eb8c0",
262 );
263 }
264
265 #[test]
266 fn it_should_not_fail_on_vec_1() {
267 check(
268 "69c2a9e62b613060c9797bdce4f6408e",
269 "4cb452d678ca946192cdc691b7ab6176",
270 "aes-cbc with iv",
271 "98ba8e075bcfa7b93abe453ab184dc68",
272 );
273 }
274
275 #[test]
276 fn it_should_not_fail_on_vec_2() {
277 check(
278 "57e79712f7b7813490b63446e1bec39f",
279 "a32fd274f2d5392f1a217aaa37a5a44d",
280 "several AES blocks means more xoring",
281 "d1aa1d92d6a93c84032ae322102aba62692b2548e92abc9f8d19bb42dd172e84c4102ff8d45889011b87de27f6d91ae4",
282 );
283 }
284}