use std::{
env,
ffi::OsString,
fs::{self, File, OpenOptions},
io::{self, BufReader, BufWriter, IsTerminal, Read, Write},
path::{Path, PathBuf},
process,
time::{Duration, Instant},
};
use zarust::{ArchiveReader, ArchiveWriter, EntryKind, NodeHandle, ROOT_NODE};
#[cfg(unix)]
use std::os::unix::ffi::{OsStrExt, OsStringExt};
const BUFFER_SIZE: usize = 64 * 1024;
fn print_help() {
println!("Usage:\n");
println!("zarust input_path [output_path]");
println!(
"If input_path is a directory, then output_path will be the ZArchive output file path"
);
println!(
"If input_path is a ZArchive file path, then output_path will be the output directory"
);
println!("output_path is optional");
}
fn main() {
process::exit(run());
}
fn run() -> i32 {
let mut args = env::args_os().skip(1);
let Some(input) = args.next() else {
print_help();
return 0;
};
let output = args.next();
if args.next().is_some() {
println!("Too many paths specified");
return -1;
}
let input = PathBuf::from(input);
if input.is_file() {
let output = output.map_or_else(|| default_extract_path(&input), PathBuf::from);
if output.exists() && !output.is_dir() {
println!("The specified output path is not a valid directory");
return -3;
}
if let Err(error) = fs::create_dir_all(&output) {
eprintln!("Failed to create output directory: {error}");
return -4;
}
extract(&input, &output)
} else if input.is_dir() {
let output = output.map_or_else(|| default_archive_path(&input), PathBuf::from);
if output.exists() && !output.is_file() {
println!("The specified output path is not a valid file");
return -10;
}
if output.exists() {
println!("The output file already exists");
return -11;
}
let result = pack(&input, &output);
if result != 0 {
let _ = fs::remove_file(&output);
}
result
} else {
println!("Input path is not a valid file or directory");
-1
}
}
fn default_extract_path(input: &Path) -> PathBuf {
let mut name = input.file_stem().unwrap_or_default().to_os_string();
name.push("_extracted");
let output = input.parent().unwrap_or_else(|| Path::new("")).join(name);
println!(
"Extracting to: {}",
output.to_string_lossy().replace('\\', "/")
);
output
}
fn default_archive_path(input: &Path) -> PathBuf {
let mut name = input.file_stem().unwrap_or_default().to_os_string();
name.push(".zar");
let output = input.parent().unwrap_or_else(|| Path::new("")).join(name);
println!(
"Outputting to: {}",
output.to_string_lossy().replace('\\', "/")
);
output
}
fn extract(input: &Path, output: &Path) -> i32 {
if !input.exists() {
println!("Unable to find archive file");
return -10;
}
let file = match File::open(input) {
Ok(file) => file,
Err(_) => {
println!("Failed to open ZArchive");
return -11;
}
};
let mut reader = match ArchiveReader::new(BufReader::new(file)) {
Ok(reader) => reader,
Err(_) => {
println!("Failed to open ZArchive");
return -11;
}
};
if let Err(error) = extract_recursive(&mut reader, ROOT_NODE, "", output) {
eprintln!("Extraction failed: {error}");
return -12;
}
0
}
fn extract_recursive<R: std::io::Read + std::io::Seek>(
reader: &mut ArchiveReader<R>,
directory: NodeHandle,
archive_path: &str,
output: &Path,
) -> zarust::Result<()> {
fs::create_dir_all(output)?;
let count = reader.directory_len(directory)?;
for index in 0..count {
let entry = reader.directory_entry(directory, index)?;
let name = safe_output_name(entry.name)?;
let kind = entry.kind;
let handle = entry.handle;
let display_name = name.to_string_lossy();
let child_archive_path = format!("{archive_path}/{display_name}");
println!("{child_archive_path}");
let child_output = output.join(&name);
match kind {
EntryKind::Directory => {
extract_recursive(reader, handle, &child_archive_path, &child_output)?;
}
EntryKind::File => extract_file(reader, handle, &child_output)?,
}
}
Ok(())
}
fn safe_output_name(name: &[u8]) -> zarust::Result<OsString> {
if name.is_empty()
|| name == b"."
|| name == b".."
|| name.contains(&b'/')
|| name.contains(&b'\\')
|| name.contains(&0)
{
return Err(zarust::Error::InvalidPath("unsafe entry name"));
}
#[cfg(unix)]
return Ok(OsString::from_vec(name.to_vec()));
#[cfg(not(unix))]
Ok(OsString::from(String::from_utf8_lossy(name).into_owned()))
}
fn extract_file<R: std::io::Read + std::io::Seek>(
reader: &mut ArchiveReader<R>,
handle: NodeHandle,
output: &Path,
) -> zarust::Result<()> {
let mut file = BufWriter::new(File::create(output)?);
let mut buffer = [0_u8; BUFFER_SIZE];
let mut offset = 0_u64;
loop {
let read = reader.read_file(handle, offset, &mut buffer)?;
if read == 0 {
break;
}
file.write_all(&buffer[..read])?;
offset += read as u64;
}
file.flush()?;
if offset != reader.file_size(handle)? {
return Err(zarust::Error::InvalidArchive(
"file data ended unexpectedly",
));
}
Ok(())
}
fn pack(input: &Path, output: &Path) -> i32 {
let mut entries = Vec::new();
let total_size = match collect_entries(input, input, output, &mut entries) {
Ok(size) => size,
Err(PackError { code, message }) => {
eprintln!("{message}");
return code;
}
};
let file = match OpenOptions::new().write(true).create_new(true).open(output) {
Ok(file) => file,
Err(_) => {
println!("Failed to create output file: {}", output.display());
return -16;
}
};
let mut writer = ArchiveWriter::new(BufWriter::new(file));
let mut progress = CompressionProgress::new(total_size);
progress.start();
match pack_entries(&mut writer, input, &entries, &mut progress) {
Ok(()) => match writer.finish() {
Ok(_) => {
progress.finish();
0
}
Err(error) => {
progress.clear();
eprintln!("Failed to finalize archive: {error}");
-16
}
},
Err(PackError { code, message }) => {
progress.clear();
eprintln!("{message}");
code
}
}
}
fn collect_entries(
root: &Path,
directory: &Path,
output: &Path,
entries: &mut Vec<PackEntry>,
) -> Result<u64, PackError> {
let directory_entries =
fs::read_dir(directory).map_err(|error| PackError::io(-15, directory, error))?;
let mut total = 0_u64;
for entry in directory_entries {
let entry = entry.map_err(|error| PackError {
code: -15,
message: format!("Failed to read input directory: {error}"),
})?;
let source = entry.path();
if same_path(&source, output) {
continue;
}
let relative = source.strip_prefix(root).unwrap();
let archive_path = archive_path(relative)?;
let file_type = entry
.file_type()
.map_err(|error| PackError::io(-15, &source, error))?;
if file_type.is_dir() {
entries.push(PackEntry {
source: source.clone(),
archive_path,
kind: PackEntryKind::Directory,
});
total = total
.checked_add(collect_entries(root, &source, output, entries)?)
.ok_or_else(PackError::too_large)?;
} else if file_type.is_file() {
let size = entry
.metadata()
.map_err(|error| PackError::io(-15, &source, error))?
.len();
total = total.checked_add(size).ok_or_else(PackError::too_large)?;
entries.push(PackEntry {
source,
archive_path,
kind: PackEntryKind::File,
});
}
}
Ok(total)
}
fn pack_entries<W: Write>(
writer: &mut ArchiveWriter<W>,
root: &Path,
entries: &[PackEntry],
progress: &mut CompressionProgress,
) -> Result<(), PackError> {
for entry in entries {
let relative = entry.source.strip_prefix(root).unwrap();
match entry.kind {
PackEntryKind::Directory => {
writer
.make_dir(&entry.archive_path, false)
.map_err(|error| PackError {
code: -13,
message: format!(
"Failed to create directory {}: {error}",
relative.display()
),
})?
}
PackEntryKind::File => {
progress.clear();
println!("Adding {}", relative.display());
progress.start();
let file = File::open(&entry.source)
.map_err(|error| PackError::io(-15, &entry.source, error))?;
add_file_with_progress(writer, &entry.archive_path, BufReader::new(file), progress)
.map_err(|error| PackError {
code: -14,
message: format!(
"Failed to create archive file {}: {error}",
relative.display()
),
})?;
}
}
}
Ok(())
}
fn add_file_with_progress<W: Write, R: Read>(
writer: &mut ArchiveWriter<W>,
archive_path: &[u8],
mut source: R,
progress: &mut CompressionProgress,
) -> zarust::Result<()> {
writer.start_file(archive_path)?;
let mut buffer = [0_u8; BUFFER_SIZE];
loop {
let read = source.read(&mut buffer)?;
if read == 0 {
break;
}
writer.append_data(&buffer[..read])?;
progress.advance(read as u64);
}
Ok(())
}
fn archive_path(path: &Path) -> Result<Vec<u8>, PackError> {
#[cfg(unix)]
{
return Ok(path.as_os_str().as_bytes().to_vec());
}
#[cfg(not(unix))]
let text = path.to_str().ok_or_else(|| PackError {
code: -14,
message: format!("Path is not valid Unicode: {}", path.display()),
})?;
#[cfg(not(unix))]
Ok(text.replace('\\', "/").into_bytes())
}
fn same_path(left: &Path, right: &Path) -> bool {
if left == right {
return true;
}
if let (Ok(left), Ok(right)) = (fs::canonicalize(left), fs::canonicalize(right)) {
return left == right;
}
let absolute = |path: &Path| {
if path.is_absolute() {
path.to_path_buf()
} else {
env::current_dir().unwrap_or_default().join(path)
}
};
absolute(left) == absolute(right)
}
struct PackError {
code: i32,
message: String,
}
struct PackEntry {
source: PathBuf,
archive_path: Vec<u8>,
kind: PackEntryKind,
}
#[derive(Clone, Copy)]
enum PackEntryKind {
Directory,
File,
}
impl PackError {
fn io(code: i32, path: &Path, error: std::io::Error) -> Self {
Self {
code,
message: format!("Failed to open input file {}: {error}", path.display()),
}
}
fn too_large() -> Self {
Self {
code: -14,
message: "Input data size exceeds the supported range".to_owned(),
}
}
}
struct CompressionProgress {
total: u64,
processed: u64,
started: Instant,
last_update: Instant,
terminal: bool,
last_width: usize,
}
impl CompressionProgress {
fn new(total: u64) -> Self {
let now = Instant::now();
Self {
total,
processed: 0,
started: now,
last_update: now,
terminal: io::stderr().is_terminal(),
last_width: 0,
}
}
fn start(&mut self) {
if self.terminal {
self.draw(false);
}
}
fn advance(&mut self, bytes: u64) {
self.processed = self.processed.saturating_add(bytes);
let interval = if self.terminal {
Duration::from_millis(100)
} else {
Duration::from_secs(1)
};
if self.last_update.elapsed() >= interval || (self.terminal && self.processed >= self.total)
{
self.draw(false);
}
}
fn finish(&mut self) {
self.draw(true);
}
fn clear(&mut self) {
if !self.terminal || self.last_width == 0 {
return;
}
let mut stderr = io::stderr().lock();
let _ = write!(stderr, "\r{:width$}\r", "", width = self.last_width);
let _ = stderr.flush();
self.last_width = 0;
}
fn draw(&mut self, finished: bool) {
let elapsed = self.started.elapsed();
let percentage = if finished || self.total == 0 {
100.0
} else {
(self.processed as f64 * 100.0 / self.total as f64).min(100.0)
};
let speed = if elapsed.is_zero() {
0.0
} else {
self.processed as f64 / elapsed.as_secs_f64()
};
let eta = if finished || self.processed >= self.total {
Some(Duration::ZERO)
} else if self.processed == 0 {
None
} else {
Some(Duration::from_secs_f64(
elapsed.as_secs_f64() * (self.total - self.processed) as f64
/ self.processed as f64,
))
};
let line = format!(
"Compressing {percentage:5.1}% | processed {} / {} | speed {}/s | elapsed {} | ETA {}",
format_bytes(self.processed as f64),
format_bytes(self.total as f64),
format_bytes(speed),
format_duration(Some(elapsed)),
format_duration(eta),
);
let mut stderr = io::stderr().lock();
if self.terminal {
let width = self.last_width.max(line.len());
let _ = write!(stderr, "\r{line:<width$}");
if finished {
let _ = writeln!(stderr);
self.last_width = 0;
} else {
self.last_width = width;
}
let _ = stderr.flush();
} else {
let _ = writeln!(stderr, "{line}");
}
self.last_update = Instant::now();
}
}
fn format_bytes(mut bytes: f64) -> String {
const UNITS: [&str; 7] = ["B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB"];
let mut unit = 0;
while bytes >= 1024.0 && unit + 1 < UNITS.len() {
bytes /= 1024.0;
unit += 1;
}
if unit == 0 {
format!("{bytes:.0} {}", UNITS[unit])
} else {
format!("{bytes:.1} {}", UNITS[unit])
}
}
fn format_duration(duration: Option<Duration>) -> String {
let Some(duration) = duration else {
return "--:--:--".to_owned();
};
let seconds = duration.as_secs();
format!(
"{:02}:{:02}:{:02}",
seconds / 3600,
seconds / 60 % 60,
seconds % 60
)
}