picca/src/main.rs

use clap::Parser;
use std::collections::{HashMap, VecDeque};
use std::fs::{self, read_to_string, File};
use std::path::PathBuf;
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, available_parallelism, JoinHandle};

mod common;
mod hashers;

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: bool,
  quiet: bool,
  print_failures: bool,
  thread_id: usize,
  filenames: Arc<Mutex<VecDeque<PathBuf>>>,
  algorithm: Arc<String>,
  hash_map: Option<Arc<Mutex<HashMap<PathBuf, String>>>>,
  file_errors: Arc<AtomicI32>,
  hash_errors: Arc<AtomicI32>,
}

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 {
      common::debug(format!(
        "thread {} is hashing file '{}'",
        info.thread_id,
        filename.as_path().display()
      ));
    }

    if filename.is_dir() {
      common::error(format!("{}: Is a directory", filename.as_path().display()));
      info.file_errors.fetch_add(1, Ordering::SeqCst);
      continue;
    }

    let file = match File::open(&filename) {
      Err(e) => {
        common::error(format!("{}: {}", filename.as_path().display(), e));
        info.file_errors.fetch_add(1, Ordering::SeqCst);
        continue;
      },
      Ok(f) => f,
    };

    let res = match &*info.algorithm.as_str() {
      "ascon" => hashers::hash_ascon(file),
      "belt" => hashers::hash_belt(file),
      "blake2b512" => hashers::hash_blake2b512(file),
      "blake2s256" => hashers::hash_blake2s256(file),
      "fsb160" => hashers::hash_fsb160(file),
      "fsb224" => hashers::hash_fsb224(file),
      "fsb256" => hashers::hash_fsb256(file),
      "fsb384" => hashers::hash_fsb384(file),
      "fsb512" => hashers::hash_fsb512(file),
      "md5" => hashers::hash_md5(file),
      "sha256" => hashers::hash_sha256(file),
      "sha384" => hashers::hash_sha384(file),
      "sha512" => hashers::hash_sha512(file),
      _ => panic!("Somehow did not pass a supported algorithm"),
    };

    match &info.hash_map {
      Some(h) => {
        if h.lock().unwrap()[&filename] == res && !info.quiet {
          println!("{}: OK", filename.as_path().display());
        } else {
          println!("{}: FAILED", filename.as_path().display());
          info.hash_errors.fetch_add(1, Ordering::SeqCst);
        }
      },
      None => {
        if !info.quiet {
          println!("{}  {}", res, filename.as_path().display());
        }
      },
    }
  }

  if !info.quiet && info.debug {
    common::debug(format!("thread {} has ran out of work", info.thread_id));
  }

  Ok(())
}

fn verify(
  cpus: usize,
  algorithm: String,
  debug: bool,
  quiet: bool,
  print_failures: bool,
  checksum_files: Vec<PathBuf>,
) -> (
  Vec<JoinHandle<Result<(), std::string::String>>>,
  Arc<AtomicI32>,
  Arc<AtomicI32>,
) {
  let mut handles = vec![];
  let mut hash_map: HashMap<PathBuf, String> = HashMap::new();
  let mut buffer = VecDeque::new();
  for file in checksum_files {
    match read_to_string(&file) {
      Err(e) => {
        common::error(format!("{}: {}", file.as_path().display(), e));
        continue;
      },
      Ok(f) => {
        for line in f.lines() {
          let split: Vec<String> = line.split_whitespace().map(|x| x.to_string()).collect();
          // println!("{}, {}", split.size_hint().0, split.size_hint().1)
          match split.len() {
            2 => {
              hash_map.insert(PathBuf::from(split[1].clone()), split[0].clone());
              buffer.push_back(PathBuf::from(split[1].clone()));
            },
            _ => common::error(format!("malformed line: {}", line)),
          }
        }
      },
    };
  }

  let arc_fe = Arc::new(AtomicI32::new(0));
  let arc_he = Arc::new(AtomicI32::new(0));
  let arc_buf = Arc::new(Mutex::new(buffer));
  let arc_hash = Arc::new(Mutex::new(hash_map));
  for i in 0..cpus {
    let safe_fe = Arc::clone(&arc_fe);
    let safe_he = Arc::clone(&arc_he);
    let safe_buf = Arc::clone(&arc_buf);
    let safe_alg = Arc::new(algorithm.clone());
    let safe_hash = Arc::clone(&arc_hash);
    handles.push(thread::spawn(move || {
      hash(ThreadInfo {
        debug,
        quiet,
        print_failures,
        thread_id: i,
        filenames: safe_buf,
        algorithm: safe_alg,
        hash_map: Some(safe_hash),
        file_errors: safe_fe,
        hash_errors: safe_he,
      })
    }))
  }

  return (handles, arc_fe, arc_he);
}

fn generate(
  cpus: usize,
  buffer: VecDeque<PathBuf>,
  algorithm: String,
  debug: bool,
  quiet: bool,
) -> (
  Vec<JoinHandle<Result<(), std::string::String>>>,
  Arc<AtomicI32>,
  Arc<AtomicI32>,
) {
  let mut handles = vec![];
  let arc_fe = Arc::new(AtomicI32::new(0));
  let arc_he = Arc::new(AtomicI32::new(0));
  let arc_buf = Arc::new(Mutex::new(buffer));
  for i in 0..cpus {
    let safe_fe = Arc::clone(&arc_fe);
    let safe_he = Arc::clone(&arc_he);
    let safe_buf = Arc::clone(&arc_buf);
    let safe_alg = Arc::new(algorithm.clone());
    handles.push(thread::spawn(move || {
      hash(ThreadInfo {
        debug,
        quiet,
        print_failures: false,
        thread_id: i,
        filenames: safe_buf,
        algorithm: safe_alg,
        hash_map: None,
        file_errors: safe_fe,
        hash_errors: safe_he,
      })
    }))
  }

  return (handles, arc_fe, arc_he);
}

fn main() {
  let args = Args::parse();
  let cpus = match args.threads {
    0 => available_parallelism().unwrap().get(),
    _ => args.threads,
  };

  if args.debug {
    common::debug(format!(
      "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 handles;
  let arc_fe;
  let arc_he;
  if &args.check.len() >= &1 {
    (handles, arc_fe, arc_he) = verify(
      cpus,
      args.algorithm,
      args.debug,
      args.quiet,
      args.failures_only,
      args.check,
    );
  } else {
    let mut buffer = VecDeque::new();
    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);
      }
    }

    (handles, arc_fe, arc_he) = generate(cpus, buffer, args.algorithm, args.debug, args.quiet);
  }

  for handle in handles {
    match handle.join().unwrap() {
      Err(e) => panic!("{}", e),
      Ok(_) => (),
    }
  }

  let fe = arc_fe.load(Ordering::SeqCst);
  let he = arc_he.load(Ordering::SeqCst);
  if fe != 0 {
    common::warning(format!("{} listed files could not be read", fe));
  }

  if he != 0 {
    common::warning(format!("{} computed checksums did NOT match", he));
  }
}