path: root/src/zip/driver.rs

use crate::driver::{ArchiveRead, ArchiveWrite, Driver};
use crate::utils::{IoCursor, ReadUtils};
use crate::zip::cp437::FromCp437;
use crate::zip::datetime::DosDateTime;
use crate::zip::structs::{
    deserialize, AesField, Cdr, Eocdr, Eocdr64, Eocdr64Locator, ExtraHeader, CDR_SIGNATURE,
    EOCDR64_LOCATOR_SIGNATURE, EOCDR64_SIGNATURE, EOCDR_SIGNATURE,
};
use crate::zip::{
    BitFlag, CompressionMethod, EncryptionMethod, ZipError, ZipFileInfo, ZipFileReader,
    ZipFileWriter, ZipResult,
};
use chrono::{DateTime, Local};
use std::collections::HashMap as Map;
use std::fs::File;
use std::io::{BufReader, Read, Seek, SeekFrom, Write};

struct Fields<'b> {
    pointer: usize,
    bytes: &'b [u8],
}

impl<'b> Fields<'b> {
    pub fn new(bytes: &'b [u8]) -> Self {
        Self { pointer: 0, bytes }
    }
}

impl<'b> Iterator for Fields<'b> {
    type Item = (u16, &'b [u8]);

    fn next(&mut self) -> Option<Self::Item> {
        let header: ExtraHeader =
            deserialize(self.bytes.get(self.pointer..self.pointer + 4)?).unwrap();
        self.pointer += 4;
        let data = self
            .bytes
            .get(self.pointer..self.pointer + header.size as usize)?;
        self.pointer += header.size as usize;
        Some((header.id, data))
    }
}

#[inline]
fn ntfs_to_local(time: u64) -> Option<DateTime<Local>> {
    Some(
        DateTime::from_timestamp(
            (time / 10000000 - 11644473600) as i64,
            (time % 10000000) as u32,
        )?
        .with_timezone(&Local),
    )
}

#[inline]
fn timestamp_to_local(time: i32) -> Option<DateTime<Local>> {
    Some(DateTime::from_timestamp(time as i64, 0)?.with_timezone(&Local))
}

pub struct Zip<Io = File> {
    io: Io,

    indexes: Map<String, usize>,
    files: Vec<ZipFileInfo>,
    comment: String,
}

impl<Io> Driver for Zip<Io> {
    type Error = ZipError;

    type Io = Io;
    type FileInfo = ZipFileInfo;
}

impl<Io: Read + Seek> ArchiveRead for Zip<Io> {
    type FileReader<'d> = ZipFileReader<'d, Io> where Io: 'd;

    fn read(mut io: Self::Io) -> ZipResult<Self> {
        // Search eocdr
        let limit = 65557.min(io.seek(SeekFrom::End(0))?) as i64;
        let start = io.seek(SeekFrom::End(-limit))?;
        let eocdr_pos = start
            + io.read_vec(
                (limit as usize)
                    .checked_sub(18)
                    .ok_or(ZipError::StructNotFound("Eocdr"))?,
            )?
            .windows(4)
            .rposition(|v| v == EOCDR_SIGNATURE)
            .ok_or(ZipError::StructNotFound("Eocdr"))? as u64;

        // Read eocdr
        io.seek(SeekFrom::Start(eocdr_pos + 4))?;
        let buf = io.read_arr::<18>()?;
        let eocdr: Eocdr = deserialize(&buf).unwrap();
        let comment = String::from_cp437(io.read_vec(eocdr.comment_len as usize)?);

        let mut cd_pointer = eocdr.cd_pointer as u64;
        let mut cd_size = eocdr.cd_size as u64;
        let mut cd_records = eocdr.cd_records as u64;

        // Try to find eocdr64locator
        if eocdr_pos >= 20 {
            io.seek(SeekFrom::Start(eocdr_pos - 20))?;
            let buf = io.read_arr::<20>()?;

            if buf[..4] == EOCDR64_LOCATOR_SIGNATURE {
                let locator: Eocdr64Locator = deserialize(&buf[4..]).unwrap();
                io.seek(SeekFrom::Start(locator.eocdr64_pointer))?;

                if io.read_arr()? != EOCDR64_SIGNATURE {
                    return Err(ZipError::InvalidSignature("Eocdr64"));
                }

                if locator.eocdr64_pointer + 76 > eocdr_pos {
                    return Err(ZipError::Overlapping("Eocdr64", "Eocdr64Locator"));
                }

                let eocdr64: Eocdr64 = deserialize(&io.read_arr::<54>()?).unwrap();
                if locator.eocdr64_pointer + eocdr64.eocdr64_size + 32 > eocdr_pos {
                    return Err(ZipError::Overlapping("Eocdr64", "Eocdr64Locator"));
                }

                cd_pointer = eocdr64.cd_pointer;
                cd_size = eocdr64.cd_size;
                cd_records = eocdr64.cd_records;

                if cd_pointer + cd_size > locator.eocdr64_pointer {
                    return Err(ZipError::Overlapping("Cdr", "Eocdr64"));
                }
            } else if cd_pointer + cd_size > eocdr_pos {
                return Err(ZipError::Overlapping("Cdr", "Eocdr"));
            }
        } else if cd_pointer + cd_size > eocdr_pos {
            return Err(ZipError::Overlapping("Cdr", "Eocdr"));
        }

        // Read cd records
        let mut indexes = Map::with_capacity(cd_records as usize);
        let mut files = Vec::with_capacity(cd_records as usize);
        io.seek(SeekFrom::Start(cd_pointer))?;
        let mut cd_reader = BufReader::with_capacity(
            cd_size.min(131072) as usize,
            IoCursor::new(&mut io, cd_pointer, cd_pointer + cd_size),
        );

        for i in 0..cd_records as usize {
            let buf = cd_reader.read_arr::<46>()?;

            if buf[..4] != CDR_SIGNATURE {
                return Err(ZipError::InvalidSignature("Cdr"));
            }
            let cdr: Cdr = deserialize(&buf[4..46]).unwrap();
            let bit_flag = BitFlag::new(cdr.bit_flag);

            let name = cd_reader.read_vec(cdr.name_len as usize)?;
            let name = if bit_flag.is_utf8() {
                String::from_utf8(name).map_err(|_| ZipError::InvalidField("file_name"))?
            } else {
                String::from_cp437(name)
            };

            let extra_fields = cd_reader.read_vec(cdr.extra_field_len as usize)?;

            let comment = cd_reader.read_vec(cdr.comment_len as usize)?;
            let comment = if bit_flag.is_utf8() {
                String::from_utf8(comment).map_err(|_| ZipError::InvalidField("file_comment"))?
            } else {
                String::from_cp437(comment)
            };

            let mut compression_method = cdr.compression_method;
            let mut encryption_method = if !bit_flag.is_encrypted() {
                EncryptionMethod::None
            } else if !bit_flag.is_strong_encryption() {
                EncryptionMethod::Weak
            } else {
                EncryptionMethod::Unsupported
            };

            let mut compressed_size = cdr.compressed_size as u64;
            let mut size = cdr.size as u64;
            let mut header_pointer = cdr.header_pointer as u64;

            let mut mtime = DateTime::from_dos_date_time(cdr.dos_date, cdr.dos_time, Local)?;
            let mut atime = None;
            let mut ctime = None;

            // Parse extensible data fields
            for (id, mut data) in Fields::new(&extra_fields) {
                match id {
                    // Zip64
                    0x0001 => {
                        if size == 0xFFFFFFFF {
                            size = u64::from_le_bytes(data.read_arr()?);
                        }
                        if compressed_size == 0xFFFFFFFF {
                            compressed_size = u64::from_le_bytes(data.read_arr()?);
                        }
                        if header_pointer == 0xFFFFFFFF {
                            header_pointer = u64::from_le_bytes(data.read_arr()?);
                        }
                    }
                    // NTFS
                    0x000a => {
                        for (id, mut data) in Fields::new(&data[4..]) {
                            match id {
                                0x0001 => {
                                    mtime = ntfs_to_local(u64::from_le_bytes(data.read_arr()?))
                                        .unwrap_or(mtime);
                                    atime = ntfs_to_local(u64::from_le_bytes(data.read_arr()?));
                                    ctime = ntfs_to_local(u64::from_le_bytes(data.read_arr()?));
                                }
                                _ => {}
                            }
                        }
                    }
                    // Unix
                    0x000d => {
                        atime = timestamp_to_local(i32::from_le_bytes(data.read_arr()?));
                        mtime = timestamp_to_local(i32::from_le_bytes(data.read_arr()?))
                            .unwrap_or(mtime);
                    }
                    // AES
                    0x9901 => {
                        let aes: AesField = deserialize(&data.read_arr::<7>()?).unwrap();
                        if aes.id != [0x41, 0x45] {
                            return Err(ZipError::InvalidField("extra_fields"));
                        }
                        encryption_method = match aes.strength {
                            0x01 => EncryptionMethod::Aes128,
                            0x02 => EncryptionMethod::Aes192,
                            0x03 => EncryptionMethod::Aes256,
                            _ => EncryptionMethod::Unsupported,
                        };
                        compression_method = aes.compression_method;
                    }
                    // Skip unrecognized header
                    _ => {}
                }
            }

            if compression_method == 99 {
                return Err(ZipError::StructNotFound("AesExtensibleData"));
            }

            indexes.insert(name.clone(), i);
            files.push(ZipFileInfo::new(
                CompressionMethod::from_struct_id(compression_method),
                encryption_method,
                bit_flag,
                mtime,
                atime,
                ctime,
                cdr.crc,
                compressed_size,
                size,
                header_pointer,
                name,
                comment,
            ));
        }

        Ok(Self {
            io,
            indexes,
            files,
            comment,
        })
    }

    fn files(&self) -> &Vec<Self::FileInfo> {
        &self.files
    }

    fn get_file_index(&self, name: &str) -> ZipResult<usize> {
        self.indexes
            .get(name)
            .ok_or(ZipError::FileNotFound)
            .copied()
    }

    fn get_file_info(&self, index: usize) -> ZipResult<&Self::FileInfo> {
        self.files.get(index).ok_or(ZipError::FileNotFound)
    }

    #[inline]
    fn get_file_reader<'d>(&'d mut self, index: usize) -> ZipResult<Self::FileReader<'d>> {
        self.get_file_reader_with_optional_password(index, None)
    }
}

impl<Io: Read + Seek> Zip<Io> {
    pub fn comment(&self) -> &String {
        &self.comment
    }

    pub fn get_file_reader_with_optional_password<'d>(
        &'d mut self,
        index: usize,
        password: Option<&[u8]>,
    ) -> ZipResult<<Self as ArchiveRead>::FileReader<'d>> {
        Ok(ZipFileReader::new(
            &mut self.io,
            self.files.get(index).ok_or(ZipError::FileNotFound)?,
            password,
        )?)
    }
}

impl<Io: Read + Write + Seek> ArchiveWrite for Zip<Io> {
    type FileWriter<'d> = ZipFileWriter<'d, Io> where Io: 'd;
}

impl<Io: Read + Write + Seek> Zip<Io> {}