Add AES encryption

8 files changed, 150 insertions, 63 deletions

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<Io> {
 }
 
 impl<Io: Seek> IoCursor<Io> {
-    pub fn new(mut io: Io, start: u64, end: u64) -> Result<Self> {
-        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<Io: Read, Aes: BlockEncrypt> {
 
     counter: u128,
     block: [u8; 16],
-    cursor: usize,
+    lower: usize,
+    upper: usize,
 }
 
 impl<Io: Read, Aes: BlockEncrypt> AesDecoder<Io, Aes> {
-    pub fn new(mut io: Io, aes: Aes) -> IoResult<Self> {
-        let block = io.read_arr::<16>()?;
+    pub fn new(io: Io, aes: Aes) -> IoResult<Self> {
         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<Io: Read, Aes: BlockEncrypt> Read for AesDecoder<Io, Aes> {
-    fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
-        todo!()
+    fn read(&mut self, mut buf: &mut [u8]) -> IoResult<usize> {
+        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<Io: Read> {
     Aes256(AesDecoder<Io, Aes256>),
 }
 
-impl<Io: Read> Encryption<Io> {
+impl<Io: Read + Seek> Encryption<IoCursor<Io>> {
     #[inline]
-    pub fn new(mut io: Io, info: &ZipFileInfo, password: Option<&[u8]>) -> ZipResult<Self> {
+    pub fn new(
+        mut io: Io,
+        cursor: u64,
+        end: u64,
+        info: &ZipFileInfo,
+        password: Option<&[u8]>,
+    ) -> ZipResult<Self> {
         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<Io: Read> Encryption<Io> {
                 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<Io: Read> Encryption<Io> {
                 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<Io: Read> Encryption<Io> {
                 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
--- /dev/null
+++ b/tests/files/archive_aes.zip
diff --git a/tests/files/archive_passwd.zip b/tests/files/archive_weak.zip
index 291bd12..291bd12 100644
--- a/tests/files/archive_passwd.zip
+++ b/tests/files/archive_weak.zip
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::<Zip>::read_from_file("tests/files/archive_passwd.zip").unwrap();
+fn test_zip_aes() {
+    let mut archive = Archive::<Zip>::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<&String>>(),
+        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::<Zip>::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();