haggis-rs/src/haggis.rs

#![warn(clippy::all, clippy::pedantic)]
use {
    clap::ArgMatches,
    haggis::{
        Algorithm, HumanSize, Listing, ListingKind, ListingStream, Message, NodeStream,
        StreamMessage,
    },
    indicatif::{ProgressBar, ProgressStyle},
    std::{
        fs::{self, File},
        io::{self, BufReader, BufWriter},
        os::fd::{AsRawFd, FromRawFd},
        process,
        sync::{Arc, atomic::{AtomicU64, Ordering}, mpsc},
        thread,
    },
    walkdir::WalkDir,
    zstd::{Decoder, Encoder},
};

mod cli;

static TEMPLATE: &str = "[ {prefix:^30!} ] {wide_bar}{pos:>5.cyan}/{len:5.green}";

fn main() {
    let matches = cli::haggis().get_matches();
    match matches.subcommand() {
        Some(("create", matches)) => {
            if let Err(e) = create(matches) {
                eprintln!("Error: {e}");
                process::exit(1);
            }
        }
        Some(("extract", matches)) => {
            if let Err(e) = extract(matches) {
                eprintln!("Error: {e}");
                process::exit(1);
            }
        }
        Some(("list", matches)) => {
            if matches.get_flag("nosort") {
                if let Err(e) = list_unsorted(matches) {
                    eprintln!("Error: {e}");
                    process::exit(1);
                }
            } else if let Err(e) = list(matches) {
                eprintln!("Error: {e}");
                process::exit(1);
            }
        }
        _ => {}
    }
}

#[allow(clippy::similar_names)]
fn create(matches: &ArgMatches) -> Result<(), haggis::Error> {
    let verbose = !matches.get_flag("stdout") || matches.get_flag("quiet");
    let algorithm: Algorithm = matches.get_one::<String>("algorithm").unwrap().parse()?;
    let uid = matches.get_one::<u32>("uid").copied();
    let gid = matches.get_one::<u32>("gid").copied();
    let mut files = vec![];
    if let Some(f) = matches.get_many::<String>("files") {
        for f in f {
            if let Ok(meta) = fs::metadata(f) {
                if meta.is_dir() {
                    let walker = WalkDir::new(f);
                    walker.into_iter().for_each(|x| {
                        if let Ok(x) = x {
                            let path = x.path().to_str().unwrap().to_string();
                            if !path.is_empty() {
                                files.push(path);
                            }
                        }
                    });
                } else {
                    files.push(f.to_string());
                }
            }
        }
    }
    let output = matches.get_one::<String>("output");
    let (sender, receiver) = mpsc::channel();
    let len = files.len();
    let mut handle = None;
    if verbose {
        let pb = ProgressBar::new(len as u64);
        pb.set_style(ProgressStyle::with_template(TEMPLATE).unwrap());
        pb.set_prefix("Adding files");
        if let Some(o) = output {
            pb.println(format!("Creating archive {o}"));
        }
        handle = Some(thread::spawn(move || {
            for msg in &receiver {
                match msg {
                    Message::NodeCreated(s) => {
                        pb.set_prefix(s.split('/').last().unwrap().to_string());
                        pb.inc(1);
                    }
                    Message::NodeSaved { name, size } => {
                        let name = name.split('/').last().unwrap();
                        pb.set_prefix(format!("{name} added, {size} bytes"));
                        pb.inc(1);
                    }
                    Message::Eof => {
                        pb.finish_and_clear();
                        break;
                    }
                    Message::Err { name, error } => {
                        pb.println(format!("Error creating node {name}: {error}"));
                    }
                }
            }
        }));
    }
    if matches.get_flag("zstd") {
        let level = matches.get_one::<i32>("level").copied().unwrap_or(3);
        if matches.get_flag("stdout") {
            let stdout = io::stdout();
            let mut writer = Encoder::new(stdout, level)?;
            haggis::par_stream_archive(&mut writer, &files, algorithm, &sender, uid, gid)?;
            let _fd = writer.finish();
        } else if let Some(o) = output {
            let fd = File::create(o)?;
            let mut writer = Encoder::new(fd, level)?;
            haggis::par_stream_archive(&mut writer, &files, algorithm, &sender, uid, gid)?;
            let _fd = writer.finish()?;
        } else {
            unreachable!();
        }
    } else if matches.get_flag("stdout") {
        let stdout = io::stdout();
        let mut writer = BufWriter::new(stdout);
        haggis::par_stream_archive(&mut writer, &files, algorithm, &sender, uid, gid)?;
    } else if let Some(o) = output {
        haggis::par_create_archive(o, &files, algorithm, &sender, uid, gid)?;
    } else {
        unreachable!();
    }
    if let Some(handle) = handle {
        match handle.join() {
            Ok(()) => {
                if verbose {
                    println!("Archive created successfully");
                }
                Ok(())
            }
            Err(e) => {
                eprintln!("Error: {e:?}");
                process::exit(1);
            }
        }
    } else {
        Ok(())
    }
}

#[allow(clippy::similar_names)]
fn extract(matches: &ArgMatches) -> Result<(), haggis::Error> {
    let total = Arc::new(AtomicU64::new(0));
    let file = matches.get_one::<String>("archive");
    let uid = matches.get_one::<u32>("uid").copied();
    let gid = matches.get_one::<u32>("gid").copied();
    let mut fd = if let Some(f) = file {
        File::open(f)?
    } else if matches.get_flag("stdin") {
        let stdin = io::stdin();
        let raw = stdin.as_raw_fd();
        unsafe { File::from_raw_fd(raw) }
    } else {
        unreachable!()
    };
    let zst = matches.get_flag("zstd")
        || if matches.get_flag("stdin") {
            false
        } else {
            haggis::detect_zstd(&mut fd)?
        };
    let dir = matches.get_one::<String>("change");
    let (sender, receiver) = mpsc::channel();
    let file = file.cloned().unwrap_or("stdin".to_string());
    let total_ref = total.clone();
    let handle = if zst {
        let reader = Decoder::new(fd)?;
        let mut stream = NodeStream::new(reader)?;
        let handle = if matches.get_flag("quiet") {
            Some(thread::spawn(move || {
                let t = progress(&file, &receiver, u64::from(stream.length));
                total_ref.store(t, Ordering::Relaxed);
                Ok::<(), haggis::Error>(())
            }))
        } else {
            None
        };
        stream.par_extract(dir.map(String::as_str), uid, gid, &sender)?;
        handle
    } else {
        let reader = BufReader::new(fd);
        let mut stream = NodeStream::new(reader)?;
        let handle = if matches.get_flag("quiet") {
            Some(thread::spawn(move || {
                let t = progress(&file, &receiver, u64::from(stream.length));
                total_ref.store(t, Ordering::Relaxed);
                Ok::<(), haggis::Error>(())
            }))
        } else {
            None
        };
        stream.par_extract(dir.map(String::as_str), uid, gid, &sender)?;
        handle
    };
    if let Some(handle) = handle {
        match handle.join() {
            Ok(_) => {
                if matches.get_flag("total") {
                    println!("{} extracted", HumanSize::from(total.load(Ordering::Relaxed)));
                } else if matches.get_flag("quiet") {
                    println!("Archive extracted successfully");
                }
                Ok(())
            }
            Err(e) => {
                eprintln!("Error: {e:?}");
                process::exit(1);
            }
        }
    } else {
        Ok(())
    }
}

fn progress(file: &str, receiver: &mpsc::Receiver<StreamMessage>, len: u64) -> u64 {
    let mut total: u64 = 0;
    let pb = ProgressBar::new(len);
    pb.set_style(ProgressStyle::with_template(TEMPLATE).unwrap());
    pb.set_prefix("Extracting files");
    pb.println(format!("Extracting archive {file}"));
    for msg in receiver {
        match msg {
            StreamMessage::FileExtracted { name, size } => {
                let name = name.split('/').last().unwrap();
                pb.set_prefix(format!("{name} extracted, {size} bytes"));
                pb.inc(1);
                total += size;
            }
            StreamMessage::LinkCreated { name, target } => {
                let name = name.split('/').last().unwrap();
                let target = target.split('/').last().unwrap();
                pb.set_prefix(format!("{name} -> {target}"));
                pb.inc(1);
            }
            StreamMessage::DirectoryCreated { name } => {
                let name = name.split('/').last().unwrap();
                pb.set_prefix(format!("mkdir {name}"));
                pb.inc(1);
            }
            StreamMessage::DeviceCreated { name } => {
                let name = name.split('/').last().unwrap();
                pb.set_prefix(format!("mknod {name}"));
                pb.inc(1);
            }
            StreamMessage::Eof => {
                pb.finish_and_clear();
                break;
            }
            StreamMessage::Err { name, error } => {
                pb.println(format!("Error with node {name}: {error}"));
            }
        }
    }
    return total;
}

fn print_listing(li: &Listing, matches: &ArgMatches) -> Result<(), haggis::Error> {
    if matches.get_flag("files") && li.kind == ListingKind::Directory {
        return Ok(());
    }
    if matches.get_flag("color") {
        if matches.get_flag("long") {
            li.print_color()?;
        } else {
            li.print_color_simple()?;
        }
    } else if matches.get_flag("long") {
        println!("{li}");
    } else {
        println!("{}", li.name);
    }
    Ok(())
}

fn list_unsorted(matches: &ArgMatches) -> Result<(), haggis::Error> {
    let file = matches.get_one::<String>("archive").unwrap();
    let fd = File::open(file)?;
    if matches.get_flag("zstd") {
        let reader = Decoder::new(fd)?;
        let stream = NodeStream::new(reader)?;
        for node in stream {
            let node = node?;
            let li = Listing::from(node);
            print_listing(&li, matches)?;
        }
    } else {
        let reader = BufReader::new(fd);
        let stream = ListingStream::new(reader)?;
        for li in stream {
            let li = li?;
            print_listing(&li, matches)?;
        }
    }
    Ok(())
}

fn list(matches: &ArgMatches) -> Result<(), haggis::Error> {
    let mut total: u64 = 0;
    let file = matches.get_one::<String>("archive").unwrap();
    let mut fd = File::open(file)?;
    let zst = matches.get_flag("zstd") || haggis::detect_zstd(&mut fd)?;
    let list = if zst {
        let reader = Decoder::new(fd)?;
        let stream = NodeStream::new(reader)?;
        let mut list = vec![];
        for node in stream {
            let node = node?;
            let listing = Listing::from(node);
            list.push(listing);
        }
        list.sort_unstable();
        list
    } else {
        let reader = BufReader::new(fd);
        let mut stream = ListingStream::new(reader)?;
        stream.list()?
    };
    for li in list {
        print_listing(&li, matches)?;
        if matches.get_flag("total") {
            if let ListingKind::Normal(s) = li.kind {
                total += s;
            }
        }
    }
    if matches.get_flag("total") {
        println!("Total: {}", HumanSize::from(total));
    }
    Ok(())
}