restic-integrity/source/main.rs

111 lines
3.5 KiB
Rust
Raw Normal View History

2022-04-10 19:49:40 +02:00
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,
2022-04-10 19:49:40 +02:00
/// 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)]
2022-04-10 19:49:40 +02:00
single_directory: bool
}
fn run() -> Result<()> {
let args = Args::parse();
let path = Path::new(&args.path);
let start = args.start;
2022-04-10 19:49:40 +02:00
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()?;
2022-04-10 19:49:40 +02:00
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(|| {
2022-04-10 19:49:40 +02:00
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);
2022-04-10 19:49:40 +02:00
}
Ok(())
});
})
2022-04-10 19:49:40 +02:00
}
fn main() {
let result = run();
if result.is_err() {
eprintln!("{}", result.unwrap_err());
std::process::exit(1);
}
}