287 lines
12 KiB
Rust
287 lines
12 KiB
Rust
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<T: Write> {
|
|
files: Vec<FileInfo>,
|
|
zip64_end_of_central_directory: Option<Zip64EndCentralDirectory>,
|
|
current_offset: u64,
|
|
data: T,
|
|
}
|
|
|
|
impl<T: Write> ZipFileWriter<T> {
|
|
/// 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<Zip64EndCentralDirectory>) -> 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<U: Read + Seek>(
|
|
&mut self,
|
|
mut file_info: FileInfo,
|
|
zip_file: &mut ZipFileReader<U>,
|
|
) {
|
|
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<U: Read + Seek>(
|
|
&mut self,
|
|
file: &mut U,
|
|
mut header: FileHeader,
|
|
local_header: Option<LocalFileHeader>,
|
|
data_descriptor: Option<DataDescriptor>,
|
|
) {
|
|
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();
|
|
}
|
|
}
|