From 4c411b76cad9cc735687dc739d2e2db5d00e5eac Mon Sep 17 00:00:00 2001
From: Igor Tolmachev <me@igorek.dev>
Date: Sun, 21 Jul 2024 16:59:14 +0900
Subject: Add AES encryption

---
 src/zip/file/read.rs | 71 +++++++++++++++++++++++++++++++++-------------------
 1 file changed, 45 insertions(+), 26 deletions(-)

(limited to 'src/zip/file')

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,
         })
     }
-- 
cgit v1.2.3