path: root/src/zip/driver.rs

use crate::driver::{ArchiveRead, ArchiveWrite, Driver};
use crate::utils::ReadUtils;
use crate::zip::cp437::FromCp437;
use crate::zip::structs::{deserialize, Cdr, Eocdr, Eocdr64, Eocdr64Locator, ExtraHeader};
use crate::zip::{
    BitFlag, CompressionMethod, EncryptionMethod, ZipError, ZipFileInfo, ZipFileReader,
    ZipFileWriter, ZipResult,
};
use chrono::{DateTime, Local, NaiveDate, NaiveDateTime, NaiveTime};
use std::collections::HashMap as Map;
use std::fs::File;
use std::io::{Read, Seek, SeekFrom, Write};

fn dos_to_local(date: u16, time: u16) -> ZipResult<DateTime<Local>> {
    Ok(NaiveDateTime::new(
        NaiveDate::from_ymd_opt(
            (date as i32 >> 9 & 0x7F) + 1980,
            date as u32 >> 5 & 0xF,
            date as u32 & 0x1F,
        )
        .ok_or(ZipError::InvalidDate)?,
        NaiveTime::from_hms_opt(
            (time as u32 >> 11) & 0x1F,
            (time as u32 >> 5) & 0x3F,
            (time as u32 & 0x1F) * 2,
        )
        .ok_or(ZipError::InvalidTime)?,
    )
    .and_local_timezone(Local)
    .unwrap())
}

fn ntfs_to_local(time: u64) -> ZipResult<DateTime<Local>> {
    Ok(DateTime::from_timestamp(
        (time / 10000000 - 11644473600) as i64,
        (time % 10000000) as u32,
    )
    .ok_or(ZipError::InvalidTime)?
    .with_timezone(&Local))
}

fn timestamp_to_local(time: i32) -> ZipResult<DateTime<Local>> {
    Ok(DateTime::from_timestamp(time as i64, 0)
        .ok_or(ZipError::InvalidTime)?
        .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 pos = start
            + io.read_vec(limit as usize - 18)?
                .windows(4)
                .rposition(|v| u32::from_le_bytes(v.try_into().unwrap()) == 0x06054b50)
                .ok_or(ZipError::EOCDRNotFound)? as u64;

        // Read eocdr
        io.seek(SeekFrom::Start(pos + 4))?;
        let buf = io.read_arr::<18>()?;
        let eocdr: Eocdr = deserialize(&buf).unwrap();
        let comment = {
            let mut buf: Vec<u8> = vec![0; eocdr.comment_len as usize];
            io.read(&mut buf)?;
            String::from_cp437(buf)
        };

        // Try to find eocdr64locator
        io.seek(SeekFrom::Start(pos - 20))?;
        let buf = io.read_arr::<20>()?;
        let (cd_pointer, cd_size, cd_records) =
            // If locator found then read eocdr64
            if u32::from_le_bytes(buf[0..4].try_into().unwrap()) == 0x07064b50 {
                let eocdr64locator: Eocdr64Locator = deserialize(&buf[4..]).unwrap();

                io.seek(SeekFrom::Start(eocdr64locator.eocdr64_pointer))?;
                let buf = io.read_arr::<56>()?;
                if u32::from_le_bytes(buf[0..4].try_into().unwrap()) != 0x06064b50 {
                    return Err(ZipError::InvalidEOCDR64Signature.into());
                }
                let eocdr64: Eocdr64 = deserialize(&buf[4..]).unwrap();

                (eocdr64.cd_pointer, eocdr64.cd_size, eocdr64.cd_records)
            } else {
                (
                    eocdr.cd_pointer as u64,
                    eocdr.cd_size as u64,
                    eocdr.cd_records as u64,
                )
            };

        // 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 buf = io.read_vec(cd_size as usize)?;

        let mut p: usize = 0;
        for i in 0..cd_records as usize {
            if u32::from_le_bytes(buf[p..p + 4].try_into().unwrap()) != 0x02014b50 {
                return Err(ZipError::InvalidCDRSignature.into());
            }
            p += 4;
            let cdr: Cdr = deserialize(&buf[p..p + 42]).unwrap();
            let bit_flag = BitFlag::new(cdr.bit_flag);
            p += 42;
            let name = if bit_flag.is_utf8() {
                String::from_utf8(buf[p..p + cdr.name_len as usize].to_vec())
                    .map_err(|_| ZipError::InvalidFileName)?
            } else {
                String::from_cp437(&buf[p..p + cdr.name_len as usize])
            };
            p += cdr.name_len as usize;
            let extra_fields: Vec<u8> = buf[p..p + cdr.extra_field_len as usize].into();
            p += cdr.extra_field_len as usize;
            let comment = if bit_flag.is_utf8() {
                String::from_utf8(buf[p..p + cdr.comment_len as usize].to_vec())
                    .map_err(|_| ZipError::InvalidFileComment)?
            } else {
                String::from_cp437(&buf[p..p + cdr.comment_len as usize])
            };
            p += cdr.comment_len as usize;

            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 = dos_to_local(cdr.dos_date, cdr.dos_time)?;
            let mut atime = None;
            let mut ctime = None;

            // Parse extensible data fields
            let mut ep: usize = 0;
            while ep < cdr.extra_field_len as usize {
                let header: ExtraHeader = deserialize(&extra_fields[ep..ep + 4]).unwrap();
                ep += 4;
                match header.id {
                    // Zip64
                    0x0001 => {
                        if size == 0xFFFFFFFF {
                            compressed_size =
                                u64::from_le_bytes(extra_fields[ep..ep + 8].try_into().unwrap());
                            ep += 8;
                        }
                        if compressed_size == 0xFFFFFFFF {
                            size = u64::from_le_bytes(extra_fields[ep..ep + 8].try_into().unwrap());
                            ep += 8;
                        }
                        if header_pointer == 0xFFFFFFFF {
                            header_pointer =
                                u64::from_le_bytes(extra_fields[ep..ep + 8].try_into().unwrap());
                            ep += 8;
                        }
                        if cdr.disk == 0xFFFF {
                            ep += 4
                        }
                    }
                    // NTFS
                    0x000a => {
                        let mut tp = ep + 4;
                        ep += header.size as usize;

                        while tp < ep {
                            let header: ExtraHeader =
                                deserialize(&extra_fields[tp..tp + 4]).unwrap();
                            tp += 4;
                            match header.id {
                                0x0001 => {
                                    mtime = ntfs_to_local(u64::from_le_bytes(
                                        extra_fields[tp..tp + 8].try_into().unwrap(),
                                    ))?;
                                    tp += 8;
                                    atime = Some(ntfs_to_local(u64::from_le_bytes(
                                        extra_fields[tp..tp + 8].try_into().unwrap(),
                                    ))?);
                                    tp += 8;
                                    ctime = Some(ntfs_to_local(u64::from_le_bytes(
                                        extra_fields[tp..tp + 8].try_into().unwrap(),
                                    ))?);
                                    tp += 8;
                                }
                                _ => {
                                    tp += header.size as usize;
                                }
                            }
                        }
                    }
                    // Unix
                    0x000d => {
                        atime = Some(timestamp_to_local(i32::from_le_bytes(
                            extra_fields[ep..ep + 4].try_into().unwrap(),
                        ))?);
                        mtime = timestamp_to_local(i32::from_le_bytes(
                            extra_fields[ep + 4..ep + 8].try_into().unwrap(),
                        ))?;
                        ep += header.size as usize
                    }
                    // Skip unrecognized header
                    _ => ep += header.size as usize,
                }
            }

            indexes.insert(name.clone(), i);
            files.push(ZipFileInfo::new(
                CompressionMethod::from_struct_id(cdr.compression_method)?,
                EncryptionMethod::from_bit_flag(bit_flag),
                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) -> crate::ArchiveResult<usize, Self::Error> {
        self.indexes
            .get(name)
            .ok_or(ZipError::FileNotFound.into())
            .copied()
    }

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

    fn get_file_reader<'d>(
        &'d mut self,
        index: usize,
        password: Option<&str>,
    ) -> ZipResult<Self::FileReader<'d>> {
        Ok(ZipFileReader::new(
            &mut self.io,
            self.files.get(index).ok_or(ZipError::FileNotFound)?,
            password,
        )?)
    }
}

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

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> {}