zarust 0.2.0

//! `zarust extract` — write an archive's contents back to the filesystem.

use std::{
    ffi::OsString,
    fs::{self, File},
    io::{Read, Seek, Write},
    path::{Path, PathBuf},
    time::Instant,
};

use zarust::{ArchiveReader, EntryKind, NodeHandle, ROOT_NODE};

use crate::cli::{
    Failure, Outcome,
    args::Options,
    ui::{Reporter, counted, display_path, format_bytes, format_duration},
};

const WRITE_BUFFER: usize = 64 * 1024;

pub fn run(
    reporter: &mut Reporter,
    archive: &Path,
    output: Option<PathBuf>,
    options: &Options,
) -> Outcome {
    let output = match output {
        Some(path) => path,
        None => default_output(archive),
    };
    if output.exists() && !output.is_dir() {
        return Err(Failure::new(format!(
            "`{}` exists and is not a directory",
            display_path(&output)
        )));
    }
    if !options.force
        && output
            .read_dir()
            .is_ok_and(|mut listing| listing.next().is_some())
    {
        return Err(Failure::new(format!(
            "`{}` is not empty (use --force to extract into it)",
            display_path(&output)
        )));
    }

    let started = Instant::now();
    let mut reader = super::open_archive(archive)?;

    if options.check {
        reporter.start_bar("checking", 0);
        let intact = reader.verify_integrity().map_err(|error| {
            Failure::new(format!("cannot read `{}`: {error}", display_path(archive)))
        })?;
        reporter.finish_bar();
        if !intact {
            return Err(Failure::new(format!(
                "`{}` failed its integrity check; refusing to extract",
                display_path(archive)
            )));
        }
    }

    // The whole tree is already in memory, so measuring it is nearly free and
    // gives the progress bar an accurate total.
    let mut totals = Totals::default();
    measure(&reader, ROOT_NODE, &mut totals)?;

    fs::create_dir_all(&output).map_err(|error| {
        Failure::new(format!(
            "cannot create `{}`: {error}",
            display_path(&output)
        ))
    })?;

    reporter.start_bar("extracting", totals.bytes);
    let mut buffer = vec![0_u8; WRITE_BUFFER];
    let result = extract_directory(reporter, &mut reader, ROOT_NODE, &output, &mut buffer);
    reporter.finish_bar();
    result?;

    let palette = reporter.palette();
    reporter.status(
        palette.green,
        "extracted",
        format_args!(
            "{bold}{}{reset}",
            display_path(&output),
            bold = palette.bold,
            reset = palette.reset
        ),
    )?;
    reporter.sub(format_args!(
        "{}, {} · {} in {}",
        counted(totals.files, "file", "files"),
        counted(totals.directories, "directory", "directories"),
        format_bytes(totals.bytes as f64),
        format_duration(Some(started.elapsed())),
    ))?;
    Ok(())
}

/// The directory `extract` uses when none is given.
///
/// Suffixed rather than bare so that packing a folder and extracting the result
/// side by side does not collide with the original: `game/` packs to `game.zar`,
/// which extracts to `game_extracted/`. This matches the original ZArchive tool.
fn default_output(archive: &Path) -> PathBuf {
    let mut name = archive.file_stem().unwrap_or_default().to_os_string();
    if name.is_empty() {
        name = OsString::from("archive");
    }
    name.push("_extracted");
    archive.parent().unwrap_or_else(|| Path::new("")).join(name)
}

#[derive(Default)]
pub struct Totals {
    pub files: u64,
    pub directories: u64,
    pub bytes: u64,
}

/// Sums the entries under `directory` without touching the compressed data.
pub fn measure<R: Read + Seek>(
    reader: &ArchiveReader<R>,
    directory: NodeHandle,
    totals: &mut Totals,
) -> Result<(), Failure> {
    let count = reader.directory_len(directory).map_err(Failure::from)?;
    for index in 0..count {
        let entry = reader
            .directory_entry(directory, index)
            .map_err(Failure::from)?;
        match entry.kind {
            EntryKind::Directory => {
                totals.directories += 1;
                measure(reader, entry.handle, totals)?;
            }
            EntryKind::File => {
                totals.files += 1;
                totals.bytes = totals.bytes.saturating_add(entry.size);
            }
        }
    }
    Ok(())
}

fn extract_directory<R: Read + Seek>(
    reporter: &mut Reporter,
    reader: &mut ArchiveReader<R>,
    directory: NodeHandle,
    output: &Path,
    buffer: &mut [u8],
) -> Result<(), Failure> {
    let count = reader.directory_len(directory).map_err(Failure::from)?;
    for index in 0..count {
        // Copy out of the borrow so the reader is free to decompress below.
        let (name, kind, handle) = {
            let entry = reader
                .directory_entry(directory, index)
                .map_err(Failure::from)?;
            (safe_name(entry.name)?, entry.kind, entry.handle)
        };
        let path = output.join(&name);
        match kind {
            EntryKind::Directory => {
                fs::create_dir_all(&path).map_err(|error| {
                    Failure::new(format!("cannot create `{}`: {error}", display_path(&path)))
                })?;
                extract_directory(reporter, reader, handle, &path, buffer)?;
            }
            EntryKind::File => {
                if reporter.is_verbose() {
                    reporter.detail(format_args!("{}", display_path(&path)))?;
                }
                extract_file(reporter, reader, handle, &path, buffer)?;
            }
        }
    }
    Ok(())
}

fn extract_file<R: Read + Seek>(
    reporter: &mut Reporter,
    reader: &mut ArchiveReader<R>,
    handle: NodeHandle,
    path: &Path,
    buffer: &mut [u8],
) -> Result<(), Failure> {
    let expected = reader.file_size(handle).map_err(Failure::from)?;
    // Data already arrives in block-sized chunks, so a BufWriter would only add
    // a copy; measured identically either way. Write straight through.
    let mut file = File::create(path).map_err(|error| {
        Failure::new(format!("cannot create `{}`: {error}", display_path(path)))
    })?;
    let mut written = 0_u64;
    loop {
        let read = reader.read_file(handle, written, buffer).map_err(|error| {
            Failure::new(format!("cannot read `{}`: {error}", display_path(path)))
        })?;
        if read == 0 {
            break;
        }
        file.write_all(&buffer[..read]).map_err(|error| {
            Failure::new(format!("cannot write `{}`: {error}", display_path(path)))
        })?;
        written += read as u64;
        reporter.advance(read as u64);
    }
    if written != expected {
        return Err(Failure::new(format!(
            "`{}` ended after {written} of {expected} bytes",
            display_path(path)
        )));
    }
    Ok(())
}

/// Rejects entry names that would escape the output directory or resolve to
/// something other than a plain child file.
fn safe_name(name: &[u8]) -> Result<OsString, Failure> {
    let unsafe_name = name.is_empty()
        || name == b"."
        || name == b".."
        || name.contains(&b'/')
        || name.contains(&b'\\')
        || name.contains(&0);
    // On Windows a name like `file:stream` opens an alternate data stream, and a
    // trailing dot or space is silently stripped, aliasing another entry.
    let unsafe_on_windows = cfg!(windows)
        && (name.contains(&b':') || name.last().is_some_and(|byte| matches!(byte, b'.' | b' ')));
    if unsafe_name || unsafe_on_windows {
        return Err(Failure::new(format!(
            "archive contains an unsafe entry name: `{}`",
            String::from_utf8_lossy(name)
        )));
    }

    #[cfg(unix)]
    {
        use std::os::unix::ffi::OsStringExt;
        Ok(OsString::from_vec(name.to_vec()))
    }
    #[cfg(not(unix))]
    Ok(OsString::from(String::from_utf8_lossy(name).into_owned()))
}

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

    #[test]
    fn rejects_traversal_and_reserved_names() {
        assert!(safe_name(b"").is_err());
        assert!(safe_name(b".").is_err());
        assert!(safe_name(b"..").is_err());
        assert!(safe_name(b"a/b").is_err());
        assert!(safe_name(b"a\\b").is_err());
        assert!(safe_name(b"a\0b").is_err());
        assert!(safe_name(b"ok.txt").is_ok());
    }

    #[test]
    #[cfg(windows)]
    fn rejects_windows_stream_and_trailing_dot_names() {
        assert!(safe_name(b"file:stream").is_err());
        assert!(safe_name(b"trailing.").is_err());
        assert!(safe_name(b"trailing ").is_err());
    }
}