use clap::Parser; use indicatif::{ParallelProgressIterator, ProgressBar, ProgressStyle, ProgressFinish}; use rayon::ThreadPoolBuilder; use rayon::iter::{ParallelIterator, IntoParallelRefIterator}; use std::fs::File; use std::path::Path; use std::io::{self, Result, Error, ErrorKind}; use sha2::{Sha256, Digest}; #[derive(Parser, Debug)] #[clap(author, version, about, long_about = None)] struct Args { /// The path to the repository path: String, /// The number of jobs to use for hashing #[clap(short, long, default_value_t = 0)] jobs: u8, #[clap(short, long, default_value_t = 0)] start: usize, /// If the path argument should be used to check all files in a directory it points to instead of a restic repository (use for testing) #[clap(long)] single_directory: bool } fn run() -> Result<()> { let args = Args::parse(); let path = Path::new(&args.path); let start = args.start; let mut directories_with_hashed_files = Vec::new(); if args.single_directory { directories_with_hashed_files.push(path.to_path_buf()); } else { directories_with_hashed_files.push(path.join(Path::new("index"))); directories_with_hashed_files.push(path.join(Path::new("keys"))); directories_with_hashed_files.push(path.join(Path::new("snapshots"))); for subdirectory in path.join(Path::new("data")).read_dir()? { let subdirectory = subdirectory?; directories_with_hashed_files.push(subdirectory.path()); } } let mut hashed_files = Vec::new(); for directory in directories_with_hashed_files { for file in directory.read_dir()? { let file = file?; let path = file.path().canonicalize()?; hashed_files.push(path); } } let progress_bar = ProgressBar::new(hashed_files.len() as u64); progress_bar.set_style(ProgressStyle::default_bar() .template("Hashing files... {bar} {pos:>7}/{len:7} [{elapsed_precise}] ") .on_finish(ProgressFinish::AndLeave)); let pool = ThreadPoolBuilder::new().num_threads(args.jobs as usize).build().unwrap(); hashed_files.sort(); let hashed_files_len = hashed_files.len(); if start != 0 { match hashed_files_len.checked_sub(start) { None => return Err(Error::new(ErrorKind::Other, format!("Specified start at {} is larger than the total number of files {}", start, hashed_files_len))), Some(remaining) => { println!("Starting at {} files, {} remaining...", start, remaining); hashed_files.drain(0..start); } } } pool.install(|| { return hashed_files.par_iter().progress_with(progress_bar).try_for_each(|path| -> Result<()> { // Just opening the file and hashing using io::copy is roughly ~2.5x fater compared to sha256::digest_file let mut hasher = Sha256::new(); let mut file = File::open(path)?; io::copy(&mut file, &mut hasher)?; let hash = format!("{:x}", hasher.finalize()); let filename = path.file_name().unwrap().to_str().unwrap(); if filename != hash { let path_string = path.to_str().unwrap(); eprintln!("Integrity check failed for {}: Expected {}, got {}", path_string, filename, hash); } Ok(()) }); }) } fn main() { let result = run(); if result.is_err() { eprintln!("{}", result.unwrap_err()); std::process::exit(1); } }