Add AES encryption

4 files changed, 85 insertions, 54 deletions

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,
         })
     }