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 ++++++++++++++++++++++++++--------------- tests/files/archive_aes.zip | Bin 0 -> 3898 bytes tests/files/archive_passwd.zip | Bin 1059 -> 0 bytes tests/files/archive_weak.zip | Bin 0 -> 1059 bytes tests/zip.rs | 67 +++++++++++++++++++++++++++++++++++--- 9 files changed, 150 insertions(+), 63 deletions(-) create mode 100644 tests/files/archive_aes.zip delete mode 100644 tests/files/archive_passwd.zip create mode 100644 tests/files/archive_weak.zip 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, }) } diff --git a/tests/files/archive_aes.zip b/tests/files/archive_aes.zip new file mode 100644 index 0000000..a64415b Binary files /dev/null and b/tests/files/archive_aes.zip differ diff --git a/tests/files/archive_passwd.zip b/tests/files/archive_passwd.zip deleted file mode 100644 index 291bd12..0000000 Binary files a/tests/files/archive_passwd.zip and /dev/null differ diff --git a/tests/files/archive_weak.zip b/tests/files/archive_weak.zip new file mode 100644 index 0000000..291bd12 Binary files /dev/null and b/tests/files/archive_weak.zip differ diff --git a/tests/zip.rs b/tests/zip.rs index 2c8fc56..9283df3 100644 --- a/tests/zip.rs +++ b/tests/zip.rs @@ -3,8 +3,60 @@ use archivator::{Archive, Zip}; use std::io::{Read, Seek, SeekFrom}; #[test] -fn test_zip_passwd() { - let mut archive = Archive::::read_from_file("tests/files/archive_passwd.zip").unwrap(); +fn test_zip_aes() { + let mut archive = Archive::::read_from_file("tests/files/archive_aes.zip").unwrap(); + + assert_eq!(archive.comment(), "archive comment"); + assert_eq!( + archive + .files() + .iter() + .map(|f| &f.name) + .collect::>(), + vec![ + "aes128/store", + "aes192/store", + "aes256/store", + "aes128/deflate", + "aes192/deflate", + "aes256/deflate", + "aes128/bzip", + "aes192/bzip", + "aes256/bzip", + "aes128/lzma", + "aes192/lzma", + "aes256/lzma", + ] + ); + + for encryption in ["aes128", "aes192", "aes256"] { + assert!(archive + .get_file_reader_by_name(&format!("{encryption}/store")) + .is_err_and(|e| e == ZipError::PasswordIsNotSpecified)); + assert!(archive + .get_file_reader_by_name_with_password("aes128/store", b"wrong_passwd") + .is_err_and(|e| e == ZipError::WrongPassword)); + + for (name, check) in [ + ("store", "98f64f03b3d168875ffa778f7fb4"), + ("deflate", "0230e7cadb76460e80cd9de611eb"), + ("bzip", "061c17646f025837e33e00425cca"), + ("lzma", "43ef5e8ed799eb7a0f25501824ff"), + ] { + let mut f = archive + .get_file_reader_by_name_with_password(&format!("{encryption}/{name}"), b"passwd") + .unwrap(); + let mut data = String::new(); + f.read_to_string(&mut data).unwrap(); + assert_eq!(&data[..24], "test encrypted file data"); + assert_eq!(&data[172..], check); + } + } +} + +#[test] +fn test_zip_weak() { + let mut archive = Archive::::read_from_file("tests/files/archive_weak.zip").unwrap(); assert_eq!(archive.comment(), "archive comment"); assert_eq!( @@ -21,9 +73,9 @@ fn test_zip_passwd() { .is_err_and(|e| e == ZipError::PasswordIsNotSpecified)); assert!(archive .get_file_reader_by_name_with_password("store", b"wrong_passwd") - .is_err_and(|e| e == ZipError::IncorrectPassword)); + .is_err_and(|e| e == ZipError::WrongPassword)); - for (name, check_data) in [ + for (name, check) in [ ("store", "1e643774f40510e37c6f3c451d9d"), ("deflate", "a70aff4b6b2754ad47852503236a"), ("bzip", "f7085f4f8ecc512a8c2c3cbe8227"), @@ -34,7 +86,7 @@ fn test_zip_passwd() { let mut data = String::new(); f.read_to_string(&mut data).unwrap(); assert_eq!(&data[..24], "test encrypted file data"); - assert_eq!(&data[172..], check_data); + assert_eq!(&data[172..], check); } } @@ -52,6 +104,11 @@ fn test_zip() { vec!["store", "deflate", "bzip", "lzma", "zstd", "xz"] ); + assert_eq!( + archive.get_file_info_by_name("none").unwrap_err(), + ZipError::FileNotFound + ); + let mut f = archive.get_file_reader_by_name("store").unwrap(); let mut data = String::new(); -- cgit v1.3