picca/src/main.rs

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<PathBuf>,

  #[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<PathBuf>,
}

struct ThreadInfo {
  debug: Arc<bool>,
  failures_only: Arc<bool>,
  quiet: Arc<bool>,
  thread_id: usize,
  filenames: Arc<Mutex<VecDeque<PathBuf>>>,
  algorithm: Arc<String>,
}

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(_) => (),
    }
  }

}