use std::io; use std::io::{Cursor, Read, Seek, SeekFrom, Write}; use crate::compression::CompressionMethod; use crate::data_descriptor::DataDescriptor; use crate::end_of_central_directory::{ EndCentralDirectory, Zip64EndCentralDirectory, Zip64EndCentralDirectoryLocator, }; use crate::{general_purpose_flags, FileHeader, FileInfo, LocalFileHeader, ZipFileReader}; use androscalpel_serializer::Serializable; use flate2::write::DeflateEncoder; use flate2::{Compression, CrcWriter}; pub struct ZipFileWriter { files: Vec, zip64_end_of_central_directory: Option, current_offset: u64, data: T, } impl ZipFileWriter { /// If provided, `zip64_info` is used to get the value that are not computable from /// the state, like version needed to extract. pub fn new(writer: T, zip64_info: Option) -> Self { Self { current_offset: 0, zip64_end_of_central_directory: zip64_info, files: vec![], data: writer, } } /// Insert a file from an existing zip pub fn insert_file_from_zip( &mut self, mut file_info: FileInfo, zip_file: &mut ZipFileReader, ) { let file_off = file_info.get_file_offset(); let file_size = file_info.get_compressed_size(); file_info.set_offset_local_header(self.current_offset); file_info.local_header.serialize(&mut self.data).unwrap(); zip_file.data.seek(SeekFrom::Start(file_off)).unwrap(); let mut remaining_size = file_size; let mut buffer = [0u8; 4096]; while remaining_size != 0 { let read = zip_file .data .read(&mut buffer[0..core::cmp::min(4096, remaining_size as usize)]) .unwrap(); self.data.write_all(&buffer[0..read]).unwrap(); remaining_size -= read as u64; } self.current_offset += file_info.local_header.size() as u64 + file_size; self.files.push(file_info); } /// Insert a file /// `header` provide for the values that are not computed from the file. /// `local_header` can be provided is information in the local header differ /// from the central directory header. pub fn insert_file( &mut self, file: &mut U, mut header: FileHeader, local_header: Option, data_descriptor: Option, ) { assert!(header.general_purpose_flags & general_purpose_flags::MASK_ENCRYPTED == 0); assert!( header.general_purpose_flags & general_purpose_flags::MASK_USE_DATA_DESCRIPTOR == 0 ); // TODO assert!(header.general_purpose_flags & general_purpose_flags::MASK_STRONG_ENCRYPTION == 0); assert!( header.general_purpose_flags & general_purpose_flags::MASK_ENCRYPTED_CENTRAL_DIR == 0 ); assert!( header.general_purpose_flags & (general_purpose_flags::MASK_COMPRESS_OPTION_1 | general_purpose_flags::MASK_COMPRESS_OPTION_2 | general_purpose_flags::MASK_UTF8_FILENAME) == 0 ); assert!( header.compression_method == CompressionMethod::Deflated || header.compression_method == CompressionMethod::Stored ); let mut local_header = if let Some(header) = local_header { assert!(header.general_purpose_flags & general_purpose_flags::MASK_ENCRYPTED == 0); assert!( header.general_purpose_flags & general_purpose_flags::MASK_USE_DATA_DESCRIPTOR == 0 ); // TODO assert!( header.general_purpose_flags & general_purpose_flags::MASK_STRONG_ENCRYPTION == 0 ); assert!( header.general_purpose_flags & general_purpose_flags::MASK_ENCRYPTED_CENTRAL_DIR == 0 ); assert!( header.general_purpose_flags & (general_purpose_flags::MASK_COMPRESS_OPTION_1 | general_purpose_flags::MASK_COMPRESS_OPTION_2 | general_purpose_flags::MASK_UTF8_FILENAME) == 0 ); assert!( header.compression_method == CompressionMethod::Deflated || header.compression_method == CompressionMethod::Stored ); LocalFileHeader { version_needed_to_extract: header.version_needed_to_extract, general_purpose_flags: if header.compression_method == CompressionMethod::Deflated { header.general_purpose_flags & (general_purpose_flags::MASK_COMPRESS_OPTION_1 | general_purpose_flags::MASK_COMPRESS_OPTION_2) } else { 0 }, compression_method: header.compression_method, last_mod_file_time: header.last_mod_file_time, last_mod_file_data: header.last_mod_file_data, crc_32: 0, compressed_size: 0, uncompressed_size: 0, file_name: header.file_name, extra_field: header.extra_field, malformed_extra_field: header.malformed_extra_field.clone(), decryption_header: None, } } else { LocalFileHeader { version_needed_to_extract: header.version_needed_to_extract, general_purpose_flags: if header.compression_method == CompressionMethod::Deflated { header.general_purpose_flags & (general_purpose_flags::MASK_COMPRESS_OPTION_1 | general_purpose_flags::MASK_COMPRESS_OPTION_2) } else { 0 }, compression_method: header.compression_method, last_mod_file_time: header.last_mod_file_time, last_mod_file_data: header.last_mod_file_data, crc_32: 0, compressed_size: 0, uncompressed_size: 0, file_name: header.file_name.clone(), extra_field: header.extra_field.clone(), malformed_extra_field: header.malformed_extra_field.clone(), decryption_header: None, } }; let pos = file.stream_position().unwrap(); let pos_end = file.seek(SeekFrom::End(0)).unwrap(); file.seek(SeekFrom::Start(pos)).unwrap(); let uncompressed_size = pos_end - pos; let header_offset = self.current_offset; local_header.set_uncompressed_size(uncompressed_size); let data = match local_header.compression_method { CompressionMethod::Stored => { let mut crc_writer = CrcWriter::new(Vec::new()); io::copy(file, &mut crc_writer).unwrap(); local_header.crc_32 = crc_writer.crc().sum(); crc_writer.into_inner() } CompressionMethod::Deflated => { // TODO: find a way to do this in place, the compressed data can be large, and storing it // in memory is bad, but Deflate consume the Reader, so we cannot juste give it self.data // TODO: Compression::default -> use flag? let mut compressor = CrcWriter::new(DeflateEncoder::new(Vec::new(), Compression::default())); io::copy(file, &mut compressor).unwrap(); local_header.crc_32 = compressor.crc().sum(); compressor.into_inner().finish().unwrap() } _ => unimplemented!(), }; local_header.set_compressed_size(data.len() as u64); local_header.serialize(&mut self.data).unwrap(); self.current_offset += local_header.size() as u64; io::copy(&mut Cursor::new(data), &mut self.data).unwrap(); self.current_offset += local_header.get_compressed_size(); header.crc_32 = local_header.crc_32; header.set_compressed_size(local_header.get_compressed_size()); header.set_uncompressed_size(local_header.get_uncompressed_size()); header.set_offset_local_header(header_offset); // TODO: compression flags are not used right now, set to zero local_header.general_purpose_flags &= !(general_purpose_flags::MASK_COMPRESS_OPTION_1 | general_purpose_flags::MASK_COMPRESS_OPTION_2); header.general_purpose_flags &= !(general_purpose_flags::MASK_COMPRESS_OPTION_1 | general_purpose_flags::MASK_COMPRESS_OPTION_2); if data_descriptor.is_some() { panic!("Writing file with data_descriptor is not yet implemented"); } let file_info = FileInfo { local_header, header, data_descriptor, }; self.files.push(file_info); } /// Finish the zip file by writing the central directory. pub fn write_central_directory(&mut self) { for file_info in &self.files { file_info.header.serialize(&mut self.data).unwrap(); } let number_of_entries_in_central_directory_on_disk = self.files.len(); let number_of_entries_in_central_directory = self.files.len(); let size_central_directory: u64 = self .files .iter() .map(|file| file.header.size() as u64) .sum(); let offset_central_directory = self.current_offset; let use_zip64 = (number_of_entries_in_central_directory_on_disk > u16::MAX as usize) || (number_of_entries_in_central_directory > u16::MAX as usize) || (size_central_directory > u32::MAX as u64) || (offset_central_directory > u32::MAX as u64) || self.zip64_end_of_central_directory.is_some(); if use_zip64 { if self.zip64_end_of_central_directory.is_none() { self.zip64_end_of_central_directory = Some(Zip64EndCentralDirectory { version_made_by: 45, version_needed_to_extract: 45, number_of_this_disk: 0, disk_number_of_central_directory_start: 0, number_entry_in_central_directory_on_this_disk: 0, number_entry_in_central_directory: 0, size_of_central_directory: 0, offset_central_directory: 0, extensible_data: vec![], }) } let zip64_edc = self.zip64_end_of_central_directory.as_mut().unwrap(); zip64_edc.number_entry_in_central_directory_on_this_disk = number_of_entries_in_central_directory_on_disk as u64; zip64_edc.number_entry_in_central_directory = number_of_entries_in_central_directory as u64; zip64_edc.size_of_central_directory = size_central_directory; zip64_edc.offset_central_directory = offset_central_directory; zip64_edc.serialize(&mut self.data).unwrap(); Zip64EndCentralDirectoryLocator { disk_number_of_zip64_end_central_directory_start: 0, offset_zip64_end_of_central_directory_record: offset_central_directory + size_central_directory, total_number_of_disks: 0, } .serialize(&mut self.data) .unwrap(); } EndCentralDirectory { disk_number: 0, disk_number_of_central_directory_start: 0, number_of_entries_in_central_directory_on_disk: if use_zip64 { u16::MAX } else { number_of_entries_in_central_directory_on_disk as u16 }, number_of_entries_in_central_directory: if use_zip64 { u16::MAX } else { number_of_entries_in_central_directory as u16 }, size_central_directory: if use_zip64 { u32::MAX } else { size_central_directory as u32 }, offset_central_directory: if use_zip64 { u32::MAX } else { offset_central_directory as u32 }, comment: vec![], } .serialize(&mut self.data) .unwrap(); } }