sevenz_rust2/
aes256sha256.rs1use std::io::{Read, Seek, Write};
2
3#[cfg(feature = "compress")]
4pub use self::enc::*;
5use crate::Password;
6use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit, generic_array::GenericArray};
7use sha2::Digest;
8
9type Aes256CbcDec = cbc::Decryptor<aes::Aes256>;
10
11#[cfg_attr(docsrs, doc(cfg(feature = "aes256")))]
12pub struct Aes256Sha256Decoder<R> {
13 cipher: Cipher,
14 input: R,
15 done: bool,
16 obuffer: Vec<u8>,
17 ostart: usize,
18 ofinish: usize,
19 pos: usize,
20}
21
22impl<R: Read> Aes256Sha256Decoder<R> {
23 pub fn new(input: R, properties: &[u8], password: &[u8]) -> Result<Self, crate::Error> {
24 let cipher = Cipher::from_properties(properties, password)?;
25 Ok(Self {
26 input,
27 cipher,
28 done: false,
29 obuffer: Default::default(),
30 ostart: 0,
31 ofinish: 0,
32 pos: 0,
33 })
34 }
35
36 fn get_more_data(&mut self) -> std::io::Result<usize> {
37 if self.done {
38 Ok(0)
39 } else {
40 self.ofinish = 0;
41 self.ostart = 0;
42 self.obuffer.clear();
43 let mut ibuffer = [0; 512];
44 let readin = self.input.read(&mut ibuffer)?;
45 if readin == 0 {
46 self.done = true;
47 self.ofinish = self.cipher.do_final(&mut self.obuffer)?;
48 Ok(self.ofinish)
49 } else {
50 let n = self
51 .cipher
52 .update(&mut ibuffer[..readin], &mut self.obuffer)?;
53 self.ofinish = n;
54 Ok(n)
55 }
56 }
57 }
58}
59
60impl<R: Read> Read for Aes256Sha256Decoder<R> {
61 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
62 if self.ostart >= self.ofinish {
63 let mut n: usize;
64 n = self.get_more_data()?;
65 while n == 0 && !self.done {
66 n = self.get_more_data()?;
67 }
68 if n == 0 {
69 return Ok(0);
70 }
71 }
72
73 if buf.is_empty() {
74 return Ok(0);
75 }
76 let buf_len = self.ofinish - self.ostart;
77 let size = buf_len.min(buf.len());
78 buf[..size].copy_from_slice(&self.obuffer[self.ostart..self.ostart + size]);
79 self.ostart += size;
80 self.pos += size;
81 Ok(size)
82 }
83}
84
85impl<R: Read + Seek> Seek for Aes256Sha256Decoder<R> {
86 fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
87 let len = self.ofinish - self.ostart;
88 match pos {
89 std::io::SeekFrom::Start(p) => {
90 let n = (p as i64 - self.pos as i64).min(len as i64);
91
92 if n < 0 {
93 Ok(0)
94 } else {
95 self.ostart += n as usize;
96 Ok(p)
97 }
98 }
99 std::io::SeekFrom::End(_) => Err(std::io::Error::new(
100 std::io::ErrorKind::Unsupported,
101 "Aes256 decoder unsupport seek from end",
102 )),
103 std::io::SeekFrom::Current(n) => {
104 let n = n.min(len as i64);
105 if n < 0 {
106 Ok(0)
107 } else {
108 self.ostart += n as usize;
109 Ok(self.pos as u64 + n as u64)
110 }
111 }
112 }
113 }
114}
115
116fn get_aes_key(properties: &[u8], password: &[u8]) -> Result<([u8; 32], [u8; 16]), crate::Error> {
117 if properties.len() < 2 {
118 return Err(crate::Error::other("AES256 properties too shart"));
119 }
120 let b0 = properties[0];
121 let num_cycles_power = b0 & 63;
122 let b1 = properties[1];
123 let iv_size = (((b0 >> 6) & 1) + (b1 & 15)) as usize;
124 let salt_size = (((b0 >> 7) & 1) + (b1 >> 4)) as usize;
125 if 2 + salt_size + iv_size > properties.len() {
126 return Err(crate::Error::other("Salt size + IV size too long"));
127 }
128 let mut salt = vec![0u8; salt_size];
129 salt.copy_from_slice(&properties[2..(2 + salt_size)]);
130 let mut iv = [0u8; 16];
131 iv[0..iv_size].copy_from_slice(&properties[(2 + salt_size)..(2 + salt_size + iv_size)]);
132 if password.is_empty() {
133 return Err(crate::Error::PasswordRequired);
134 }
135 let aes_key = if num_cycles_power == 0x3f {
136 let mut aes_key = [0u8; 32];
137 aes_key.copy_from_slice(&salt[..salt_size]);
138 let n = password.len().min(aes_key.len() - salt_size);
139 aes_key[salt_size..n + salt_size].copy_from_slice(&password[0..n]);
140 aes_key
141 } else {
142 let mut sha = sha2::Sha256::default();
143 let mut extra = [0u8; 8];
144 for _ in 0..(1u32 << num_cycles_power) {
145 sha.update(&salt);
146 sha.update(password);
147 sha.update(extra);
148 for item in &mut extra {
149 *item = item.wrapping_add(1);
150 if *item != 0 {
151 break;
152 }
153 }
154 }
155 sha.finalize().into()
156 };
157 Ok((aes_key, iv))
158}
159
160struct Cipher {
161 dec: Aes256CbcDec,
162 buf: Vec<u8>,
163}
164
165impl Cipher {
166 fn from_properties(properties: &[u8], password: &[u8]) -> Result<Self, crate::Error> {
167 let (aes_key, iv) = get_aes_key(properties, password)?;
168 Ok(Self {
169 dec: Aes256CbcDec::new(&GenericArray::from(aes_key), &iv.into()),
170 buf: Default::default(),
171 })
172 }
173
174 fn update<W: Write>(&mut self, mut data: &mut [u8], mut output: W) -> std::io::Result<usize> {
175 let mut n = 0;
176 if !self.buf.is_empty() {
177 assert!(self.buf.len() < 16);
178 let end = 16 - self.buf.len();
179 self.buf.extend_from_slice(&data[..end]);
180 data = &mut data[end..];
181 let block = GenericArray::from_mut_slice(&mut self.buf);
182 self.dec.decrypt_block_mut(block);
183 let out = block.as_slice();
184 output.write_all(out)?;
185 n += out.len();
186 self.buf.clear();
187 }
188
189 for a in data.chunks_mut(16) {
190 if a.len() < 16 {
191 self.buf.extend_from_slice(a);
192 break;
193 }
194 let block = GenericArray::from_mut_slice(a);
195 self.dec.decrypt_block_mut(block);
196 let out = block.as_slice();
197 output.write_all(out)?;
198 n += out.len();
199 }
200 Ok(n)
201 }
202
203 fn do_final(&mut self, output: &mut Vec<u8>) -> std::io::Result<usize> {
204 if self.buf.is_empty() {
205 output.clear();
206 Ok(0)
207 } else {
208 Err(std::io::Error::new(
209 std::io::ErrorKind::InvalidData,
210 "IllegalBlockSize",
211 ))
212 }
213 }
214}
215#[cfg(feature = "compress")]
216mod enc {
217 type Aes256CbcEnc = cbc::Encryptor<aes::Aes256>;
218
219 use super::*;
220
221 #[cfg_attr(docsrs, doc(cfg(feature = "aes256")))]
222 pub struct Aes256Sha256Encoder<W> {
223 output: W,
224 enc: Aes256CbcEnc,
225 buffer: Vec<u8>,
226 finished: bool,
227 write_size: u32,
228 }
229
230 #[cfg_attr(docsrs, doc(cfg(feature = "aes256")))]
231 #[derive(Debug, Clone)]
232 pub struct AesEncoderOptions {
233 pub password: Password,
234 pub iv: [u8; 16],
235 pub salt: [u8; 16],
236 pub num_cycles_power: u8,
237 }
238
239 impl AesEncoderOptions {
240 pub fn new(password: Password) -> Self {
241 let mut iv = [0; 16];
242 getrandom::fill(&mut iv).expect("Can't generate IV");
243
244 let mut salt = [0; 16];
245 getrandom::fill(&mut salt).expect("Can't generate salt");
246
247 Self {
248 password,
249 iv,
250 salt,
251 num_cycles_power: 8,
252 }
253 }
254
255 pub fn properties(&self) -> [u8; 34] {
256 let mut props = [0u8; 34];
257 self.write_properties(&mut props);
258 props
259 }
260
261 #[inline]
262 pub fn write_properties(&self, props: &mut [u8]) {
263 assert!(props.len() >= 34);
264 props[0] = (self.num_cycles_power & 0x3f) | 0xc0;
265 props[1] = 0xff;
266 props[2..18].copy_from_slice(&self.salt);
267 props[18..34].copy_from_slice(&self.iv);
268 }
269 }
270
271 impl<W> Aes256Sha256Encoder<W> {
272 pub fn new(output: W, options: &AesEncoderOptions) -> Result<Self, crate::Error> {
273 let (key, iv) = get_aes_key(&options.properties(), options.password.as_slice())?;
274
275 Ok(Self {
276 output,
277 enc: Aes256CbcEnc::new(&GenericArray::from(key), &iv.into()),
278 buffer: Default::default(),
279 finished: false,
280 write_size: 0,
281 })
282 }
283
284 #[inline(always)]
285 fn write_block(&mut self, block: &mut [u8]) -> std::io::Result<()>
286 where
287 W: Write,
288 {
289 let block2 = GenericArray::from_mut_slice(block);
290 self.enc.encrypt_block_mut(block2);
291 self.output.write_all(block)?;
292 self.write_size += block.len() as u32;
293 Ok(())
294 }
295 }
296
297 impl<W: Write> Write for Aes256Sha256Encoder<W> {
298 fn write(&mut self, mut buf: &[u8]) -> std::io::Result<usize> {
299 if self.finished && !buf.is_empty() {
300 return Ok(0);
301 }
302 if buf.is_empty() {
303 self.finished = true;
304 self.flush()?;
305 return self.output.write(buf);
306 }
307 let len = buf.len();
308 if !self.buffer.is_empty() {
309 assert!(self.buffer.len() < 16);
310 if buf.len() + self.buffer.len() >= 16 {
311 let buffer = &self.buffer[..];
312 let end = 16 - buffer.len();
313
314 let mut block = [0u8; 16];
315 block[0..buffer.len()].copy_from_slice(buffer);
316 block[buffer.len()..16].copy_from_slice(&buf[..end]);
317 self.write_block(&mut block)?;
318 self.buffer.clear();
319 buf = &buf[end..];
320 } else {
321 self.buffer.extend_from_slice(buf);
322 return Ok(len);
323 }
324 }
325
326 for data in buf.chunks(16) {
327 if data.len() < 16 {
328 self.buffer.extend_from_slice(data);
329 break;
330 }
331 let mut block = [0u8; 16];
332 block.copy_from_slice(data);
333 self.write_block(&mut block)?;
334 }
335
336 Ok(len)
337 }
338
339 fn flush(&mut self) -> std::io::Result<()> {
340 if !self.buffer.is_empty() && self.finished {
341 assert!(self.buffer.len() < 16);
342 let mut block = [0u8; 16];
343 block[..self.buffer.len()].copy_from_slice(&self.buffer);
344 self.write_block(&mut block)?;
345 self.buffer.clear();
346 }
347 Ok(())
348 }
349 }
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355 use std::io::Cursor;
356
357 #[allow(clippy::unused_io_amount)]
358 #[cfg(feature = "compress")]
359 #[test]
360 fn test_aes_codec() {
361 let mut encoded = vec![];
362 let writer = Cursor::new(&mut encoded);
363 let pwd: Password = "1234".into();
364 let options = AesEncoderOptions::new(pwd.clone());
365 let mut enc = Aes256Sha256Encoder::new(writer, &options).unwrap();
366 let original = include_bytes!("./aes256sha256.rs");
367 enc.write_all(original).expect("encode data");
368 enc.write(&[]).unwrap();
369
370 let mut encoded_data = &encoded[..];
371 let mut dec =
372 Aes256Sha256Decoder::new(&mut encoded_data, &options.properties(), pwd.as_slice())
373 .unwrap();
374
375 let mut decoded = vec![];
376 std::io::copy(&mut dec, &mut decoded).unwrap();
377 assert_eq!(&decoded[..original.len()], &original[..]);
378 }
379}