From 7bcdc3b4ca460aec2b98fb2dca6165788c562b05 Mon Sep 17 00:00:00 2001
From: Igor Tolmachev <me@igorek.dev>
Date: Sat, 20 Jul 2024 16:52:39 +0900
Subject: Partial aes implementation and others improvements

---
 src/zip/file/read.rs | 99 +++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 87 insertions(+), 12 deletions(-)

(limited to 'src/zip/file/read.rs')

diff --git a/src/zip/file/read.rs b/src/zip/file/read.rs
index c26b304..80bdcfb 100644
--- a/src/zip/file/read.rs
+++ b/src/zip/file/read.rs
@@ -1,11 +1,16 @@
 use crate::driver::FileDriver;
 use crate::utils::{IoCursor, ReadUtils};
-use crate::zip::encryption::WeakDecoder;
+use crate::zip::encryption::{AesDecoder, Keys, WeakDecoder};
+use crate::zip::structs::FILE_HEADER_SIGNATURE;
 use crate::zip::{CompressionMethod, EncryptionMethod, ZipError, ZipFileInfo, ZipResult};
+use aes::cipher::KeyInit;
+use aes::{Aes128, Aes192, Aes256};
 use bzip2::read::BzDecoder;
 use flate2::read::DeflateDecoder;
 use liblzma::read::XzDecoder;
 use liblzma::stream::{Filters, LzmaOptions, Stream};
+use pbkdf2::pbkdf2_hmac_array;
+use sha1::Sha1;
 use std::io::{
     BufReader, Error as IoError, ErrorKind as IoErrorKind, Read, Result as IoResult, Seek, SeekFrom,
 };
@@ -14,34 +19,101 @@ use zstd::stream::Decoder as ZstdDecoder;
 enum Encryption<Io: Read> {
     None(Io),
     Weak(WeakDecoder<Io>),
+    Aes128(AesDecoder<Io, Aes128>),
+    Aes192(AesDecoder<Io, Aes192>),
+    Aes256(AesDecoder<Io, Aes256>),
 }
 
 impl<Io: Read> Encryption<Io> {
-    pub fn new(io: Io, info: &ZipFileInfo, password: Option<&[u8]>) -> ZipResult<Self> {
+    #[inline]
+    pub fn new(mut io: Io, info: &ZipFileInfo, password: Option<&[u8]>) -> ZipResult<Self> {
         Ok(match info.encryption_method {
             EncryptionMethod::None => Self::None(io),
-            EncryptionMethod::Weak(check) => Self::Weak(WeakDecoder::new(
-                io,
-                check,
-                password.ok_or(ZipError::PasswordIsNotSpecified)?,
-            )?),
-            EncryptionMethod::Unsupported => {
-                return Err(ZipError::UnsupportedEncryptionMethod.into())
+            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);
+                }
+
+                Self::Weak(WeakDecoder::new(io, keys))
+            }
+            EncryptionMethod::Aes128 => {
+                let header = io.read_arr::<10>()?;
+                let salt = &header[..8];
+                let check = &header[8..];
+
+                let hash = pbkdf2_hmac_array::<Sha1, 34>(
+                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
+                    salt,
+                    1000,
+                );
+                let key = &hash[..16];
+
+                if check != &hash[32..] {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Aes128(AesDecoder::new(io, Aes128::new(key.into()))?)
+            }
+            EncryptionMethod::Aes192 => {
+                let header = io.read_arr::<14>()?;
+                let salt = &header[..12];
+                let check = &header[12..];
+
+                let hash = pbkdf2_hmac_array::<Sha1, 50>(
+                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
+                    salt,
+                    1000,
+                );
+                let key = &hash[..24];
+
+                if check != &hash[48..] {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Aes192(AesDecoder::new(io, Aes192::new(key.into()))?)
+            }
+            EncryptionMethod::Aes256 => {
+                let header = io.read_arr::<18>()?;
+                let salt = &header[..16];
+                let check = &header[16..];
+
+                let hash = pbkdf2_hmac_array::<Sha1, 66>(
+                    password.ok_or(ZipError::PasswordIsNotSpecified)?,
+                    salt,
+                    1000,
+                );
+                let key = &hash[..32];
+
+                if check != &hash[64..] {
+                    return Err(ZipError::IncorrectPassword);
+                }
+
+                Self::Aes256(AesDecoder::new(io, Aes256::new(key.into()))?)
             }
+            EncryptionMethod::Unsupported => return Err(ZipError::UnsupportedEncryptionMethod),
         })
     }
 }
 
 impl<Io: Read> Read for Encryption<Io> {
+    #[inline]
     fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
         match self {
             Self::None(io) => io.read(buf),
             Self::Weak(io) => io.read(buf),
+            Self::Aes128(io) => io.read(buf),
+            Self::Aes192(io) => io.read(buf),
+            Self::Aes256(io) => io.read(buf),
         }
     }
 }
 
 impl<Io: Read + Seek> Seek for Encryption<Io> {
+    #[inline]
     fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
         match self {
             Self::None(io) => io.seek(pos),
@@ -62,6 +134,7 @@ enum Compression<Io: Read> {
 }
 
 impl<Io: Read + Seek> Compression<Io> {
+    #[inline]
     pub fn new(mut io: Io, info: &ZipFileInfo) -> ZipResult<Self> {
         Ok(match info.compression_method {
             CompressionMethod::Store => Self::Store(io),
@@ -88,13 +161,14 @@ impl<Io: Read + Seek> Compression<Io> {
             CompressionMethod::Zstd => Self::Zstd(ZstdDecoder::new(io)?),
             CompressionMethod::Xz => Self::Xz(XzDecoder::new(io)),
             CompressionMethod::Unsupported(id) => {
-                return Err(ZipError::UnsupportedCompressionMethod(id).into())
+                return Err(ZipError::UnsupportedCompressionMethod(id))
             }
         })
     }
 }
 
 impl<Io: Read> Read for Compression<Io> {
+    #[inline]
     fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
         match self {
             Compression::Store(io) => io.read(buf),
@@ -107,6 +181,7 @@ impl<Io: Read> Read for Compression<Io> {
 }
 
 impl<Io: Read + Seek> Seek for Compression<Io> {
+    #[inline]
     fn seek(&mut self, pos: SeekFrom) -> IoResult<u64> {
         match self {
             Compression::Store(io) => io.seek(pos),
@@ -137,8 +212,8 @@ impl<'d, Io: Read + Seek> ZipFileReader<'d, Io> {
         io.seek(SeekFrom::Start(info.header_pointer))?;
 
         let buf = io.read_arr::<30>()?;
-        if u32::from_le_bytes(buf[..4].try_into().unwrap()) != 0x04034b50 {
-            return Err(ZipError::InvalidFileHeaderSignature.into());
+        if buf[..4] != FILE_HEADER_SIGNATURE {
+            return Err(ZipError::InvalidFileHeaderSignature);
         }
         let data_pointer = info.header_pointer
             + 30
-- 
cgit v1.2.3