use clap::Parser; use std::collections::VecDeque; use std::fs::{self, File}; use std::path::PathBuf; use std::sync::{Arc, Mutex}; use std::thread::{self, available_parallelism}; mod hashers; mod common; const ALGORITHMS: [&'static str; 3] = ["sha256", "sha384", "sha512"]; const UNSECURE_ALGORITHMS: [&'static str; 1] = ["md5"]; #[derive(Parser)] #[command(name = "psha")] #[command(version = option_env!("CARGO_PKG_VERSION"))] #[command(about = "A parallel checksum tool for various algorithms")] #[command(long_about = None)] struct Args { #[arg( short, long, help = "Use at most this number of threads, 0 means as many as there are processor cores", default_value = "0" )] threads: usize, #[arg( short, long, help = "Enable debug output (thread info, algorithm used/detected)" )] debug: bool, #[arg( short, long, help = "Specify an algorithm for hashing", default_value = "sha256", value_parser = { let mut cleaned: Vec<&str> = vec![]; for i in ALGORITHMS { cleaned.push(i); } for i in UNSECURE_ALGORITHMS { cleaned.push(i); } clap::builder::PossibleValuesParser::new(Vec::from(cleaned)) } )] algorithm: String, #[arg( short = 'f', long, help = "Show canonicalized (relative paths converted to absolute) file paths" )] canonicalize: bool, #[arg( short = 'c', long, help = "Read checksums from the file(s) and verify them" )] check: Vec, #[arg( short = 'q', long, help = "(only used with -c) Only print checksums that fail; do not print OK for files that are successful" )] failures_only: bool, #[arg( short = 'Q', long, help = "(only used with -c) Suppress all output to stdout, including failures" )] quiet: bool, #[arg(trailing_var_arg = true)] files: Vec, } struct ThreadInfo { debug: Arc, failures_only: Arc, quiet: Arc, thread_id: usize, filenames: Arc>>, algorithm: Arc, } fn hash(info: ThreadInfo) -> Result<(), String> { loop { let filename = match info.filenames.lock().unwrap().pop_front() { Some(f) => f, None => break, }; if !*info.quiet && *info.debug { eprintln!( "thread {} is hashing file '{}'", info.thread_id, filename.as_path().display() ); } let file = File::open(&filename).unwrap(); let res = match &*info.algorithm.as_str() { "sha256" => hashers::hash_sha256(file), "sha384" => hashers::hash_sha384(file), "sha512" => hashers::hash_sha512(file), "md5" => hashers::hash_md5(file), _ => panic!("Somehow did not pass a supported algorithm"), }; if !*info.quiet { println!("{} {}", res, filename.as_path().display()); } } if !*info.quiet && *info.debug { eprintln!("thread {} has ran out of work", info.thread_id); } Ok(()) } fn main() { let args = Args::parse(); let mut buffer = VecDeque::new(); let mut handles = vec![]; for file in args.files { if args.canonicalize { match fs::canonicalize(file.as_path()) { Ok(p) => buffer.push_back(p), Err(e) => panic!("unable to canonicalize {}: {}", file.as_path().display(), e), }; } else { buffer.push_back(file); } } let cpus = match args.threads { 0 => available_parallelism().unwrap().get(), _ => args.threads }; if args.debug { eprintln!("Starting psha using algorithm {} with {} threads", args.algorithm, cpus) } if UNSECURE_ALGORITHMS.contains(&args.algorithm.as_str()) { common::warning(format!("{} is an unsecure hashing algorithm!", &args.algorithm)); } let arc_buf = Arc::new(Mutex::new(buffer)); for i in 0..cpus { let safe_buf = Arc::clone(&arc_buf); let safe_alg = Arc::new(args.algorithm.clone()); handles.push(thread::spawn(move || { hash(ThreadInfo { debug: Arc::new(args.debug), failures_only: Arc::new(args.failures_only), quiet: Arc::new(args.quiet), thread_id: i, filenames: safe_buf, algorithm: safe_alg, }) })) } for handle in handles { match handle.join().unwrap() { Err(e) => panic!("{}", e), Ok(_) => (), } } }