office_rs/common/

zip_utils.rs

//! ZIP文件处理工具模块
//! 提供Office文档ZIP压缩包操作的通用功能

use crate::context::ErrorContext;
use crate::error::{ OfficeError, Result };
use std::collections::HashMap;
use std::fs::File;
use std::io::{ BufReader, Cursor, Read, Seek, Write };
use std::path::{ Component, Path };
use time::OffsetDateTime;
use zip::write::FileOptions;
use zip::{ CompressionMethod, ZipArchive };

/// ZIP操作安全配置
#[derive(Debug, Clone)]
pub struct ZipSecurityConfig {
    /// 最大解压缩大小（字节）
    pub max_uncompressed_size: u64,
    /// 最大单个文件大小（字节）
    pub max_file_size: u64,
    /// 最大文件数量
    pub max_file_count: usize,
    /// 是否允许路径遍历
    pub allow_path_traversal: bool,
    /// 内存缓冲区大小（字节）
    pub memory_buffer_size: u64,
}

impl Default for ZipSecurityConfig {
    fn default() -> Self {
        Self {
            max_uncompressed_size: 100 * 1024 * 1024, // 100MB
            max_file_size: 50 * 1024 * 1024, // 50MB
            max_file_count: 1000, // 1000个文件
            allow_path_traversal: false,
            memory_buffer_size: 50 * 1024 * 1024, // 50MB缓冲区
        }
    }
}

impl ZipSecurityConfig {
    /// 创建宽松的安全配置（用于可信文件）
    pub fn permissive() -> Self {
        Self {
            max_uncompressed_size: 1024 * 1024 * 1024, // 1GB
            max_file_size: 500 * 1024 * 1024, // 500MB
            max_file_count: 10000,
            allow_path_traversal: false,
            memory_buffer_size: 500 * 1024 * 1024, // 500MB
        }
    }

    /// 创建严格的安全配置（用于不可信文件）
    pub fn strict() -> Self {
        Self {
            max_uncompressed_size: 10 * 1024 * 1024, // 10MB
            max_file_size: 5 * 1024 * 1024, // 5MB
            max_file_count: 100,
            allow_path_traversal: false,
            memory_buffer_size: 5 * 1024 * 1024, // 5MB
        }
    }
}

/// ZIP文档条目信息
#[derive(Debug, Clone)]
pub struct ZipEntry {
    pub name: String,
    pub size: u64,
    pub compressed_size: u64,
    pub is_directory: bool,
    pub last_modified: Option<std::time::SystemTime>,
}

/// 路径安全验证函数
fn validate_zip_path(path: &str, allow_traversal: bool) -> Result<()> {
    if !allow_traversal {
        let path_obj = Path::new(path);
        for component in path_obj.components() {
            match component {
                Component::ParentDir => {
                    return Err(OfficeError::Other(format!("检测到路径遍历攻击: {}", path)));
                }
                Component::RootDir => {
                    return Err(OfficeError::Other(format!("检测到绝对路径: {}", path)));
                }
                _ => {}
            }
        }
    }
    Ok(())
}

impl ZipEntry {
    /// 创建新的ZIP条目
    pub fn new(name: String) -> Self {
        Self {
            name,
            size: 0,
            compressed_size: 0,
            is_directory: false,
            last_modified: None,
        }
    }

    /// 检查是否为文件
    pub fn is_file(&self) -> bool {
        !self.is_directory
    }

    /// 获取文件扩展名
    pub fn extension(&self) -> Option<&str> {
        Path::new(&self.name)
            .extension()
            .and_then(|ext| ext.to_str())
    }

    /// 获取文件名（不含路径）
    pub fn file_name(&self) -> Option<&str> {
        Path::new(&self.name)
            .file_name()
            .and_then(|name| name.to_str())
    }

    /// 获取目录路径
    pub fn parent_path(&self) -> Option<&str> {
        Path::new(&self.name)
            .parent()
            .and_then(|path| path.to_str())
    }
}

/// ZIP文档读取器
pub struct ZipReader<R: Read + Seek> {
    archive: ZipArchive<R>,
    entries: HashMap<String, ZipEntry>,
    security_config: ZipSecurityConfig,
    total_uncompressed_size: u64,
}

impl ZipReader<BufReader<File>> {
    /// 从文件路径打开ZIP文档（使用默认安全配置）
    pub fn open_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        Self::open_file_with_config(path, ZipSecurityConfig::default())
    }

    /// 从文件路径打开ZIP文档（使用自定义安全配置）
    pub fn open_file_with_config<P: AsRef<Path>>(
        path: P,
        config: ZipSecurityConfig
    ) -> Result<Self> {
        let file = File::open(&path).map_err(|_e| {
            OfficeError::file_not_found_with_context(
                path.as_ref().to_string_lossy().to_string(),
                ErrorContext {
                    operation: Some("打开ZIP文件".to_string()),
                    ..Default::default()
                }
            )
        })?;

        let reader = BufReader::new(file);
        Self::new_with_config(reader, config)
    }
}

impl<R: Read + Seek> ZipReader<R> {
    /// 创建新的ZIP读取器（使用默认安全配置）
    pub fn new(reader: R) -> Result<Self> {
        Self::new_with_config(reader, ZipSecurityConfig::default())
    }

    /// 创建新的ZIP读取器（使用自定义安全配置）
    pub fn new_with_config(reader: R, config: ZipSecurityConfig) -> Result<Self> {
        let mut archive = ZipArchive::new(reader).map_err(|e| {
            OfficeError::Zip(e).with_context(ErrorContext {
                operation: Some("创建ZIP读取器".to_string()),
                ..Default::default()
            })
        })?;

        // 检查文件数量限制
        if archive.len() > config.max_file_count {
            return Err(
                OfficeError::Other(
                    format!("ZIP文件包含过多文件: {} > {}", archive.len(), config.max_file_count)
                )
            );
        }

        let mut entries = HashMap::new();
        let mut total_uncompressed_size = 0u64;

        // 读取所有条目信息
        for i in 0..archive.len() {
            let file = archive.by_index(i).map_err(|e| {
                OfficeError::Zip(e).with_context(ErrorContext {
                    operation: Some("读取ZIP条目".to_string()),
                    ..Default::default()
                })
            })?;

            let file_name = file.name();

            // 验证路径安全性
            validate_zip_path(file_name, config.allow_path_traversal)?;

            // 检查单个文件大小限制
            if file.size() > config.max_file_size {
                return Err(
                    OfficeError::Other(
                        format!(
                            "文件过大: {} ({} 字节) > {} 字节",
                            file_name,
                            file.size(),
                            config.max_file_size
                        )
                    )
                );
            }

            total_uncompressed_size = total_uncompressed_size.saturating_add(file.size());

            let mut entry = ZipEntry::new(file_name.to_string());
            entry.size = file.size();
            entry.compressed_size = file.compressed_size();
            entry.is_directory = file.is_dir();
            entry.last_modified = file.last_modified().and_then(|dt| {
                OffsetDateTime::try_from(dt)
                    .ok()
                    .map(|offset_dt| {
                        std::time::SystemTime::UNIX_EPOCH +
                            std::time::Duration::from_secs(offset_dt.unix_timestamp() as u64)
                    })
            });

            entries.insert(file_name.to_string(), entry);
        }

        // 检查总解压缩大小限制
        if total_uncompressed_size > config.max_uncompressed_size {
            return Err(
                OfficeError::Other(
                    format!(
                        "ZIP文件解压缩后过大: {} 字节 > {} 字节",
                        total_uncompressed_size,
                        config.max_uncompressed_size
                    )
                )
            );
        }

        Ok(Self {
            archive,
            entries,
            security_config: config,
            total_uncompressed_size,
        })
    }

    /// 获取所有条目列表
    pub fn entries(&self) -> &HashMap<String, ZipEntry> {
        &self.entries
    }

    /// 检查文件是否存在
    pub fn contains_file(&self, name: &str) -> bool {
        self.entries.contains_key(name)
    }

    /// 获取文件条目信息
    pub fn get_entry(&self, name: &str) -> Option<&ZipEntry> {
        self.entries.get(name)
    }

    /// 读取文件内容为字节数组
    pub fn read_file(&mut self, name: &str) -> Result<Vec<u8>> {
        let mut file = self.archive.by_name(name).map_err(|e| {
            OfficeError::Zip(e).with_context(ErrorContext {
                operation: Some(format!("读取ZIP文件: {}", name)),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        // 检查文件大小是否超过内存缓冲区限制
        if file.size() > self.security_config.memory_buffer_size {
            return Err(
                OfficeError::Other(
                    format!(
                        "文件过大，无法加载到内存: {} ({} 字节) > {} 字节",
                        name,
                        file.size(),
                        self.security_config.memory_buffer_size
                    )
                )
            );
        }

        // 安全地创建缓冲区，避免整数溢出
        let size = file.size() as usize;
        if size > (isize::MAX as usize) {
            return Err(OfficeError::Other(format!("文件大小超出系统限制: {} 字节", size)));
        }

        let mut contents = Vec::with_capacity(size);
        file.read_to_end(&mut contents).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some(format!("读取文件内容: {}", name)),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        Ok(contents)
    }

    /// 读取文件内容为字符串
    pub fn read_file_to_string(&mut self, name: &str) -> Result<String> {
        let bytes = self.read_file(name)?;
        String::from_utf8(bytes).map_err(|e| OfficeError::Other(format!("UTF-8解码错误: {}", e)))
    }

    /// 提取文件到指定路径
    pub fn extract_file<P: AsRef<Path>>(&mut self, name: &str, output_path: P) -> Result<()> {
        // 验证输出路径安全性
        let output_path_str = output_path.as_ref().to_string_lossy();
        validate_zip_path(&output_path_str, self.security_config.allow_path_traversal)?;

        let mut file = self.archive.by_name(name).map_err(|e| {
            OfficeError::Zip(e).with_context(ErrorContext {
                operation: Some(format!("提取ZIP文件: {}", name)),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        // 检查文件大小限制
        if file.size() > self.security_config.max_file_size {
            return Err(
                OfficeError::Other(
                    format!(
                        "文件过大，无法提取: {} ({} 字节) > {} 字节",
                        name,
                        file.size(),
                        self.security_config.max_file_size
                    )
                )
            );
        }

        let mut output_file = File::create(&output_path).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some("创建输出文件".to_string()),
                file_path: Some(output_path.as_ref().to_string_lossy().to_string()),
                ..Default::default()
            })
        })?;

        std::io::copy(&mut file, &mut output_file).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some("复制文件内容".to_string()),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        Ok(())
    }

    /// 提取所有文件到指定目录
    pub fn extract_all<P: AsRef<Path>>(&mut self, output_dir: P) -> Result<()> {
        let output_dir = output_dir.as_ref();
        let mut total_extracted_size = 0u64;

        // 使用HashSet避免重复文件名收集的性能问题
        let mut processed_files = std::collections::HashSet::new();

        // 先收集所有需要提取的文件名
        let file_names: Vec<String> = self.entries
            .iter()
            .filter(|(_, entry)| !entry.is_directory)
            .map(|(name, _)| name.clone())
            .collect();

        for name in file_names {
            // 避免重复处理
            if !processed_files.insert(name.clone()) {
                continue;
            }

            let entry = &self.entries[&name];

            // 检查累计提取大小
            total_extracted_size = total_extracted_size.saturating_add(entry.size);
            if total_extracted_size > self.security_config.max_uncompressed_size {
                return Err(
                    OfficeError::Other(
                        format!(
                            "提取的文件总大小超过限制: {} 字节 > {} 字节",
                            total_extracted_size,
                            self.security_config.max_uncompressed_size
                        )
                    )
                );
            }

            let output_path = output_dir.join(&name);

            // 验证输出路径安全性
            let output_path_str = output_path.to_string_lossy();
            validate_zip_path(&output_path_str, self.security_config.allow_path_traversal)?;

            // 创建父目录
            if let Some(parent) = output_path.parent() {
                std::fs::create_dir_all(parent).map_err(|e| {
                    OfficeError::Io(e).with_context(ErrorContext {
                        operation: Some("创建目录".to_string()),
                        file_path: Some(parent.to_string_lossy().to_string()),
                        ..Default::default()
                    })
                })?;
            }

            self.extract_file(&name, &output_path)?;
        }

        Ok(())
    }

    /// 列出指定目录下的文件
    pub fn list_files_in_directory(&self, dir_path: &str) -> Vec<&ZipEntry> {
        let normalized_dir = if dir_path.is_empty() {
            "".to_string()
        } else if dir_path.ends_with('/') {
            dir_path.to_string()
        } else {
            format!("{}/", dir_path)
        };

        self.entries
            .values()
            .filter(|entry| {
                entry.name.starts_with(&normalized_dir) &&
                    entry.name != normalized_dir &&
                    !entry.name[normalized_dir.len()..].contains('/')
            })
            .collect()
    }

    /// 查找匹配模式的文件
    pub fn find_files(&self, pattern: &str) -> Vec<&ZipEntry> {
        self.entries
            .values()
            .filter(|entry| entry.name.contains(pattern))
            .collect()
    }
}

/// ZIP文档写入器
pub struct ZipWriter<W: Write + Seek> {
    writer: zip::ZipWriter<W>,
    written_files: Vec<String>,
}

impl ZipWriter<File> {
    /// 创建新的ZIP文件
    pub fn create_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file = File::create(&path).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some("创建ZIP文件".to_string()),
                file_path: Some(path.as_ref().to_string_lossy().to_string()),
                ..Default::default()
            })
        })?;

        Self::new(file)
    }
}

impl<W: Write + Seek> ZipWriter<W> {
    /// 创建新的ZIP写入器
    pub fn new(writer: W) -> Result<Self> {
        let zip_writer = zip::ZipWriter::new(writer);
        Ok(Self {
            writer: zip_writer,
            written_files: Vec::new(),
        })
    }

    /// 添加文件到ZIP
    pub fn add_file(&mut self, name: &str, data: &[u8]) -> Result<()> {
        let options = FileOptions::<()>
            ::default()
            .compression_method(CompressionMethod::Deflated)
            .unix_permissions(0o644);

        self.writer.start_file(name, options).map_err(|e| {
            OfficeError::Zip(e).with_context(ErrorContext {
                operation: Some(format!("开始写入文件: {}", name)),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        self.writer.write_all(data).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some(format!("写入文件数据: {}", name)),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        self.written_files.push(name.to_string());
        Ok(())
    }

    /// 添加字符串文件到ZIP
    pub fn add_file_from_string(&mut self, name: &str, content: &str) -> Result<()> {
        self.add_file(name, content.as_bytes())
    }

    /// 添加目录到ZIP
    pub fn add_directory(&mut self, name: &str) -> Result<()> {
        let dir_name = if name.ends_with('/') { name.to_string() } else { format!("{}/", name) };

        let options = FileOptions::<()>::default().compression_method(CompressionMethod::Stored);

        self.writer.start_file(&dir_name, options).map_err(|e| {
            OfficeError::Zip(e).with_context(ErrorContext {
                operation: Some(format!("创建目录: {}", name)),
                file_path: Some(name.to_string()),
                ..Default::default()
            })
        })?;

        self.written_files.push(dir_name);
        Ok(())
    }

    /// 从现有文件添加到ZIP
    pub fn add_file_from_path<P: AsRef<Path>>(
        &mut self,
        zip_path: &str,
        file_path: P
    ) -> Result<()> {
        let mut file = File::open(&file_path).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some("打开源文件".to_string()),
                file_path: Some(file_path.as_ref().to_string_lossy().to_string()),
                ..Default::default()
            })
        })?;

        let mut buffer = Vec::new();
        file.read_to_end(&mut buffer).map_err(|e| {
            OfficeError::Io(e).with_context(ErrorContext {
                operation: Some("读取源文件".to_string()),
                file_path: Some(file_path.as_ref().to_string_lossy().to_string()),
                ..Default::default()
            })
        })?;

        self.add_file(zip_path, &buffer)
    }

    /// 获取已写入的文件列表
    pub fn written_files(&self) -> &[String] {
        &self.written_files
    }

    /// 完成ZIP文件写入
    pub fn finish(self) -> Result<W> {
        self.writer.finish().map_err(|e| {
            OfficeError::Zip(e).with_context(ErrorContext {
                operation: Some("完成ZIP文件写入".to_string()),
                ..Default::default()
            })
        })
    }
}

/// ZIP工具函数
pub mod utils {
    use super::*;

    /// 检查文件是否为ZIP格式
    pub fn is_zip_file<P: AsRef<Path>>(path: P) -> bool {
        if let Ok(file) = File::open(path) {
            let reader = BufReader::new(file);
            ZipArchive::new(reader).is_ok()
        } else {
            false
        }
    }

    /// 获取ZIP文件信息
    pub fn get_zip_info<P: AsRef<Path>>(path: P) -> Result<(usize, u64, u64)> {
        let reader = ZipReader::open_file(path)?;
        let entries = reader.entries();

        let file_count = entries.len();
        let total_size = entries
            .values()
            .map(|e| e.size)
            .sum();
        let total_compressed_size = entries
            .values()
            .map(|e| e.compressed_size)
            .sum();

        Ok((file_count, total_size, total_compressed_size))
    }

    /// 验证ZIP文件完整性
    pub fn validate_zip<P: AsRef<Path>>(path: P) -> Result<bool> {
        let mut reader = ZipReader::open_file(path)?;

        // 收集所有文件名
        let file_names: Vec<String> = reader
            .entries()
            .iter()
            .filter(|(_, entry)| entry.is_file())
            .map(|(name, _)| name.clone())
            .collect();

        // 尝试读取所有文件
        for name in file_names {
            let _data = reader.read_file(&name)?;
            // 如果能成功读取，说明文件完整
        }

        Ok(true)
    }

    /// 创建内存中的ZIP
    pub fn create_memory_zip(files: &[(String, Vec<u8>)]) -> Result<Vec<u8>> {
        let mut buffer = Vec::new();
        {
            let cursor = Cursor::new(&mut buffer);
            let mut writer = ZipWriter::new(cursor)?;

            for (name, data) in files {
                writer.add_file(name, data)?;
            }

            writer.finish()?;
        }

        Ok(buffer)
    }

    /// 从内存中读取ZIP
    pub fn read_memory_zip(data: &[u8]) -> Result<ZipReader<Cursor<&[u8]>>> {
        let cursor = Cursor::new(data);
        ZipReader::new(cursor)
    }

    /// 复制ZIP文件中的特定文件到新ZIP
    pub fn copy_zip_files<P1: AsRef<Path>, P2: AsRef<Path>>(
        source_path: P1,
        target_path: P2,
        file_patterns: &[&str]
    ) -> Result<()> {
        let mut source_reader = ZipReader::open_file(source_path)?;
        let mut target_writer = ZipWriter::create_file(target_path)?;

        // 收集匹配的文件名
        let mut files_to_copy = Vec::new();
        for pattern in file_patterns {
            let matching_files = source_reader.find_files(pattern);
            for entry in matching_files {
                if entry.is_file() {
                    files_to_copy.push(entry.name.clone());
                }
            }
        }

        // 复制文件
        for file_name in files_to_copy {
            let data = source_reader.read_file(&file_name)?;
            target_writer.add_file(&file_name, &data)?;
        }

        target_writer.finish()?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Cursor;

    #[test]
    fn test_memory_zip_creation() {
        let files = vec![
            ("test1.txt".to_string(), b"Hello World".to_vec()),
            ("test2.txt".to_string(), b"Goodbye World".to_vec())
        ];

        let zip_data = utils::create_memory_zip(&files).unwrap();
        assert!(!zip_data.is_empty());

        let mut reader = utils::read_memory_zip(&zip_data).unwrap();
        assert!(reader.contains_file("test1.txt"));
        assert!(reader.contains_file("test2.txt"));

        let content1 = reader.read_file_to_string("test1.txt").unwrap();
        assert_eq!(content1, "Hello World");
    }

    #[test]
    fn test_zip_entry() {
        let mut entry = ZipEntry::new("folder/test.xml".to_string());
        entry.size = 1024;

        assert_eq!(entry.file_name(), Some("test.xml"));
        assert_eq!(entry.extension(), Some("xml"));
        assert_eq!(entry.parent_path(), Some("folder"));
        assert!(entry.is_file());
    }
}