From 4c411b76cad9cc735687dc739d2e2db5d00e5eac Mon Sep 17 00:00:00 2001 From: Igor Tolmachev Date: Sun, 21 Jul 2024 16:59:14 +0900 Subject: Add AES encryption --- src/utils/cursor.rs | 7 ++--- src/zip/encryption/aes.rs | 49 ++++++++++++++++++++++++-------- src/zip/encryption/weak.rs | 15 +--------- src/zip/error.rs | 4 +-- src/zip/file/read.rs | 71 +++++++++++++++++++++++++++++----------------- 5 files changed, 88 insertions(+), 58 deletions(-) (limited to 'src') diff --git a/src/utils/cursor.rs b/src/utils/cursor.rs index 91d44f8..ee80474 100644 --- a/src/utils/cursor.rs +++ b/src/utils/cursor.rs @@ -7,13 +7,12 @@ pub struct IoCursor { } impl IoCursor { - pub fn new(mut io: Io, start: u64, end: u64) -> Result { - let cursor = io.seek(SeekFrom::Start(start))?; - Ok(Self { + pub fn new(io: Io, cursor: u64, end: u64) -> Self { + Self { io, cursor, bounds: (cursor, end), - }) + } } } diff --git a/src/zip/encryption/aes.rs b/src/zip/encryption/aes.rs index 6f41aaa..b690482 100644 --- a/src/zip/encryption/aes.rs +++ b/src/zip/encryption/aes.rs @@ -1,4 +1,3 @@ -use crate::utils::ReadUtils; use aes::cipher::generic_array::GenericArray; use aes::cipher::BlockEncrypt; use std::io::{Read, Result as IoResult}; @@ -10,37 +9,63 @@ pub struct AesDecoder { counter: u128, block: [u8; 16], - cursor: usize, + lower: usize, + upper: usize, } impl AesDecoder { - pub fn new(mut io: Io, aes: Aes) -> IoResult { - let block = io.read_arr::<16>()?; + pub fn new(io: Io, aes: Aes) -> IoResult { let mut decoder = Self { io, aes, - counter: 1, - block, - cursor: 0, + + counter: 0, + block: [0; 16], + lower: 0, + upper: 0, }; - decoder.decrypt_block(); + + decoder.update_block()?; + Ok(decoder) } #[inline] - fn decrypt_block(&mut self) { + fn update_block(&mut self) -> IoResult<()> { + self.upper = self.io.read(&mut self.block)?; + self.lower = 0; + + self.counter += 1; let mut mask = self.counter.to_le_bytes(); + self.aes .encrypt_block(GenericArray::from_mut_slice(&mut mask)); + for (b, m) in self.block.iter_mut().zip(mask) { *b ^= m } - self.counter += 1; + + Ok(()) } } impl Read for AesDecoder { - fn read(&mut self, buf: &mut [u8]) -> IoResult { - todo!() + fn read(&mut self, mut buf: &mut [u8]) -> IoResult { + let mut bytes = 0; + while !buf.is_empty() && self.lower != self.upper { + for (to, fr) in buf.iter_mut().zip(&self.block[self.lower..self.upper]) { + *to = *fr + } + + let consumed = buf.len().min(self.upper - self.lower); + buf = &mut buf[consumed..]; + self.lower += consumed; + bytes += consumed; + + if self.lower == 16 { + self.update_block()?; + } + } + Ok(bytes) } } diff --git a/src/zip/encryption/weak.rs b/src/zip/encryption/weak.rs index 144cd53..ebddb2d 100644 --- a/src/zip/encryption/weak.rs +++ b/src/zip/encryption/weak.rs @@ -45,7 +45,7 @@ impl Keys { } } - fn update(&mut self, byte: u8) { + pub fn update(&mut self, byte: u8) { self.key0 = crc32(byte, self.key0); self.key1 = self .key1 @@ -55,19 +55,6 @@ impl Keys { self.key2 = crc32((self.key1 >> 24) as u8, self.key2); } - pub fn set_password(&mut self, passwd: &[u8]) { - for b in passwd { - self.update(*b) - } - } - - pub fn set_header(&mut self, header: [u8; 12]) -> u8 { - for b in &header[..11] { - self.decode_byte(*b); - } - self.decode_byte(header[11]) - } - #[allow(dead_code)] pub fn encode_bytes(&mut self, byte: u8) -> u8 { let key = self.key2 | 2; diff --git a/src/zip/error.rs b/src/zip/error.rs index 5508177..87cd9e9 100644 --- a/src/zip/error.rs +++ b/src/zip/error.rs @@ -23,7 +23,7 @@ pub enum ZipError { InvalidFileComment, FileNotFound, - IncorrectPassword, + WrongPassword, PasswordIsNotSpecified, CompressedDataIsUnseekable, EncryptedDataIsUnseekable, @@ -81,7 +81,7 @@ impl Display for ZipError { Self::InvalidFileComment => write!(f, "Invalid file comment"), Self::FileNotFound => write!(f, "File not found"), - Self::IncorrectPassword => write!(f, "Incorrect password"), + Self::WrongPassword => write!(f, "Wrong password"), Self::PasswordIsNotSpecified => write!(f, "Password is not specified"), Self::CompressedDataIsUnseekable => write!(f, "Compressed data is unseekable"), Self::EncryptedDataIsUnseekable => write!(f, "Encrypted data is unseekable"), diff --git a/src/zip/file/read.rs b/src/zip/file/read.rs index 80bdcfb..d25655e 100644 --- a/src/zip/file/read.rs +++ b/src/zip/file/read.rs @@ -24,21 +24,34 @@ enum Encryption { Aes256(AesDecoder), } -impl Encryption { +impl Encryption> { #[inline] - pub fn new(mut io: Io, info: &ZipFileInfo, password: Option<&[u8]>) -> ZipResult { + pub fn new( + mut io: Io, + cursor: u64, + end: u64, + info: &ZipFileInfo, + password: Option<&[u8]>, + ) -> ZipResult { Ok(match info.encryption_method { - EncryptionMethod::None => Self::None(io), + EncryptionMethod::None => Self::None(IoCursor::new(io, cursor, end)), EncryptionMethod::Weak => { let mut keys = Keys::new(); - keys.set_password(password.ok_or(ZipError::PasswordIsNotSpecified)?); - let check = keys.set_header(io.read_arr()?); - if check != info.password_check() { - return Err(ZipError::IncorrectPassword); + for b in password.ok_or(ZipError::PasswordIsNotSpecified)? { + keys.update(*b) + } + + let header = io.read_arr::<12>()?; + for b in &header[..11] { + keys.decode_byte(*b); } - Self::Weak(WeakDecoder::new(io, keys)) + if keys.decode_byte(header[11]) != info.password_check() { + return Err(ZipError::WrongPassword); + } + + Self::Weak(WeakDecoder::new(IoCursor::new(io, cursor + 12, end), keys)) } EncryptionMethod::Aes128 => { let header = io.read_arr::<10>()?; @@ -53,10 +66,13 @@ impl Encryption { let key = &hash[..16]; if check != &hash[32..] { - return Err(ZipError::IncorrectPassword); + return Err(ZipError::WrongPassword); } - Self::Aes128(AesDecoder::new(io, Aes128::new(key.into()))?) + Self::Aes128(AesDecoder::new( + IoCursor::new(io, cursor + 10, end - 10), + Aes128::new(key.into()), + )?) } EncryptionMethod::Aes192 => { let header = io.read_arr::<14>()?; @@ -71,10 +87,13 @@ impl Encryption { let key = &hash[..24]; if check != &hash[48..] { - return Err(ZipError::IncorrectPassword); + return Err(ZipError::WrongPassword); } - Self::Aes192(AesDecoder::new(io, Aes192::new(key.into()))?) + Self::Aes192(AesDecoder::new( + IoCursor::new(io, cursor + 14, end - 10), + Aes192::new(key.into()), + )?) } EncryptionMethod::Aes256 => { let header = io.read_arr::<18>()?; @@ -89,10 +108,13 @@ impl Encryption { let key = &hash[..32]; if check != &hash[64..] { - return Err(ZipError::IncorrectPassword); + return Err(ZipError::WrongPassword); } - Self::Aes256(AesDecoder::new(io, Aes256::new(key.into()))?) + Self::Aes256(AesDecoder::new( + IoCursor::new(io, cursor + 18, end - 10), + Aes256::new(key.into()), + )?) } EncryptionMethod::Unsupported => return Err(ZipError::UnsupportedEncryptionMethod), }) @@ -215,20 +237,17 @@ impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> { if buf[..4] != FILE_HEADER_SIGNATURE { return Err(ZipError::InvalidFileHeaderSignature); } - let data_pointer = info.header_pointer - + 30 - + u16::from_le_bytes(buf[26..28].try_into().unwrap()) as u64 - + u16::from_le_bytes(buf[28..30].try_into().unwrap()) as u64; + + let cursor = io.seek(SeekFrom::Start( + info.header_pointer + + 30 + + u16::from_le_bytes(buf[26..28].try_into().unwrap()) as u64 + + u16::from_le_bytes(buf[28..30].try_into().unwrap()) as u64, + ))?; + let end = cursor + info.compressed_size; Ok(Self { - io: Compression::new( - Encryption::new( - IoCursor::new(io, data_pointer, data_pointer + info.compressed_size)?, - info, - password, - )?, - info, - )?, + io: Compression::new(Encryption::new(io, cursor, end, info, password)?, info)?, info, }) } -- cgit v1.2.3